<commit_before>package tokens

import (
	"github.com/dgrijalva/jwt-go"
	"github.com/go-errors/errors"
	"github.com/leanovate/microzon-auth-go/config"
	"github.com/leanovate/microzon-auth-go/logging"
	"github.com/leanovate/microzon-auth-go/revocations"
	"github.com/leanovate/microzon-auth-go/store"
	"time"
)

type TokenManager struct {
	logger logging.Logger
	config *config.TokenConfig
	store  store.Store
}

func NewTokenManager(config *config.TokenConfig, store store.Store, parent logging.Logger) *TokenManager {

	return &TokenManager{
		logger: parent.WithContext(map[string]interface{}{"package": "tokens"}),
		config: config,
		store:  store,
	}
}

func (t *TokenManager) CreateToken(realm, subject string) (*TokenInfoVO, error) {
	certWithKey, err := t.store.SelfCertificate()
	if err != nil {
		return nil, err
	}
	expiresAt := time.Now().Add(time.Duration(t.config.TokenTTL) * time.Second)
	return newTokenInfo(realm, subject, expiresAt, certWithKey)
}

func (t *TokenManager) RefreshToken(token *jwt.Token) (interface{}, error) {
	certWithKey, err := t.store.SelfCertificate()
	if err != nil {
		return nil, err
	}
	expiresAt := time.Now().Add(t.config.TokenTTL)
	return refreshToken(token, expiresAt, certWithKey)
}

func (t *TokenManager) RevokeToken(token *jwt.Token) (interface{}, error) {
	sha256 := revocations.RawSha256FromData(token.Raw)
	expiresAt := time.Unix((int64)(token.Claims[jwtClaimExpiresAt].(float64)), 0)
	return nil, t.store.AddRevocation(sha256, expiresAt)
}

func (r *TokenManager) TokenHandler(token *jwt.Token) (interface{}, error) {
	r.logger.Debugf("Token: %v", token)
	err := r.verifyRevocations(token)
	if err != nil {
		return nil, err
	} else {
		return r.verifyCertificate(token)
	}
}

func (r *TokenManager) verifyCertificate(token *jwt.Token) (interface{}, error) {
	x5t := token.Header[jwtHeaderThumbprint]
	if x5t != nil {
		if cert, err := r.store.CertificateByThumbprint(x5t.(string)); err != nil || cert == nil {
			return nil, errors.Errorf("Certificate not found: %s", x5t)
		} else if cert.NotAfter.Before(time.Now()) {
			return nil, errors.Errorf("Certificate is expired: %s", x5t)
		} else {
			return cert.PublicKey, nil
		}
	} else {
		return nil, errors.New("No x5t field found in token")
	}
}

func (r *TokenManager) verifyRevocations(token *jwt.Token) error {
	sha256 := revocations.RawSha256FromData(token.Raw)
	revoked, err := r.store.IsRevoked(sha256)
	if err != nil {
		return err
	}
	if revoked {
		return errors.Errorf("Token has been revoked: %s", sha256.String())
	}
	return nil
}
<commit_msg>Fix revocation ttl<commit_after>package tokens

import (
	"github.com/dgrijalva/jwt-go"
	"github.com/go-errors/errors"
	"github.com/leanovate/microzon-auth-go/config"
	"github.com/leanovate/microzon-auth-go/logging"
	"github.com/leanovate/microzon-auth-go/revocations"
	"github.com/leanovate/microzon-auth-go/store"
	"time"
)

type TokenManager struct {
	logger logging.Logger
	config *config.TokenConfig
	store  store.Store
}

func NewTokenManager(config *config.TokenConfig, store store.Store, parent logging.Logger) *TokenManager {

	return &TokenManager{
		logger: parent.WithContext(map[string]interface{}{"package": "tokens"}),
		config: config,
		store:  store,
	}
}

func (t *TokenManager) CreateToken(realm, subject string) (*TokenInfoVO, error) {
	certWithKey, err := t.store.SelfCertificate()
	if err != nil {
		return nil, err
	}
	expiresAt := time.Now().Add(t.config.TokenTTL * time.Second)
	return newTokenInfo(realm, subject, expiresAt, certWithKey)
}

func (t *TokenManager) RefreshToken(token *jwt.Token) (interface{}, error) {
	certWithKey, err := t.store.SelfCertificate()
	if err != nil {
		return nil, err
	}
	expiresAt := time.Now().Add(t.config.TokenTTL)
	return refreshToken(token, expiresAt, certWithKey)
}

func (t *TokenManager) RevokeToken(token *jwt.Token) (interface{}, error) {
	sha256 := revocations.RawSha256FromData(token.Raw)
	expiresAt := time.Unix((int64)(token.Claims[jwtClaimExpiresAt].(float64)), 0)
	return nil, t.store.AddRevocation(sha256, expiresAt)
}

func (r *TokenManager) TokenHandler(token *jwt.Token) (interface{}, error) {
	r.logger.Debugf("Token: %v", token)
	err := r.verifyRevocations(token)
	if err != nil {
		return nil, err
	} else {
		return r.verifyCertificate(token)
	}
}

func (r *TokenManager) verifyCertificate(token *jwt.Token) (interface{}, error) {
	x5t := token.Header[jwtHeaderThumbprint]
	if x5t != nil {
		if cert, err := r.store.CertificateByThumbprint(x5t.(string)); err != nil || cert == nil {
			return nil, errors.Errorf("Certificate not found: %s", x5t)
		} else if cert.NotAfter.Before(time.Now()) {
			return nil, errors.Errorf("Certificate is expired: %s", x5t)
		} else {
			return cert.PublicKey, nil
		}
	} else {
		return nil, errors.New("No x5t field found in token")
	}
}

func (r *TokenManager) verifyRevocations(token *jwt.Token) error {
	sha256 := revocations.RawSha256FromData(token.Raw)
	revoked, err := r.store.IsRevoked(sha256)
	if err != nil {
		return err
	}
	if revoked {
		return errors.Errorf("Token has been revoked: %s", sha256.String())
	}
	return nil
}
<|endoftext|>


<commit_before>package python

import (
	"fmt"
	"log"
	"regexp"
	"strings"

	"sourcegraph.com/sourcegraph/doc"
)

var rstDirectiveRegex = regexp.MustCompile(`\.\. [a-zA-Z]+:: `)

// Properly format python documentation by doing the following:
// Check if txt is rST. If so, compile into html, otherwise return html approximation to plaintext.
func formatDocs(txt string) string {
	var format doc.Formatter
	if strings.Contains(txt, "~~~") || len(rstDirectiveRegex.FindStringIndex(txt)) > 0 {
		format = doc.ReStructuredText

		var err error
		txt, err = rstPreprocess(txt)
		if err != nil {
			log.Printf("rstPreprocess failed: %s", err)
		}
	} else if strings.Contains(txt, "====") {
		format = doc.Markdown
	} else {
		format = doc.Text
	}

	html, err := doc.ToHTML(format, []byte(txt))
	if err != nil {
		log.Printf("doc.ToHTML failed: %s", err)
	}
	return string(html)
}

var indentMatcher = regexp.MustCompile("(\\s*)[^\\s].*")
var rstRoleMatcher = regexp.MustCompile(":[A-Za-z]+:`([^`]+)`")

// Preprocess docstring before it can be fed to rst2html.  This is
// necessary, because sometimes indentations in docstrings are
// formatted to look good with code, rather than be valid rst.
func rstPreprocess(txt string) (string, error) {
	lines := strings.Split(txt, "\n")

	globalIndent := ""
	foundFirstIndent := false
	for l, line := range lines {
		if l == 0 && !strings.HasPrefix(line, "..") {
			continue
		} else if strings.TrimSpace(line) == "" {
			continue
		}

		matches := indentMatcher.FindStringSubmatch(line)
		if matches == nil || len(matches) < 2 {
			return "", fmt.Errorf("Could not process line: %s", line)
		}

		if !foundFirstIndent {
			foundFirstIndent = true
			globalIndent = matches[1]
		} else {
			if len(globalIndent) > len(matches[1]) {
				globalIndent = matches[1]
			}
		}
	}

	for l, line := range lines {
		if strings.HasPrefix(line, globalIndent) {
			lines[l] = line[len(globalIndent):]
		}
	}

	indentNormalizedTxt := strings.Join(lines, "\n")
	roleRemovedTxt := rstRoleMatcher.ReplaceAllString(indentNormalizedTxt, ":code:`$1`")

	return roleRemovedTxt, nil
}
<commit_msg>handle possible unprintable characters in docstrings<commit_after>package python

import (
	"fmt"
	"log"
	"regexp"
	"strings"

	"sourcegraph.com/sourcegraph/doc"
)

var rstDirectiveRegex = regexp.MustCompile(`\.\. [a-zA-Z]+:: `)

// Properly format python documentation by doing the following:
// Check if txt is rST. If so, compile into html, otherwise return html approximation to plaintext.
func formatDocs(txt string) string {
	txt = doc.EscapeUnprintable(txt)

	var format doc.Formatter
	if strings.Contains(txt, "~~~") || len(rstDirectiveRegex.FindStringIndex(txt)) > 0 {
		format = doc.ReStructuredText

		var err error
		txt, err = rstPreprocess(txt)
		if err != nil {
			log.Printf("rstPreprocess failed: %s", err)
		}
	} else if strings.Contains(txt, "====") {
		format = doc.Markdown
	} else {
		format = doc.Text
	}

	html, err := doc.ToHTML(format, []byte(txt))
	if err != nil {
		log.Printf("doc.ToHTML failed: %s", err)
	}
	return string(html)
}

var indentMatcher = regexp.MustCompile("(\\s*)[^\\s].*")
var rstRoleMatcher = regexp.MustCompile(":[A-Za-z]+:`([^`]+)`")

// Preprocess docstring before it can be fed to rst2html.  This is
// necessary, because sometimes indentations in docstrings are
// formatted to look good with code, rather than be valid rst.
func rstPreprocess(txt string) (string, error) {
	lines := strings.Split(txt, "\n")

	globalIndent := ""
	foundFirstIndent := false
	for l, line := range lines {
		if l == 0 && !strings.HasPrefix(line, "..") {
			continue
		} else if strings.TrimSpace(line) == "" {
			continue
		}

		matches := indentMatcher.FindStringSubmatch(line)
		if matches == nil || len(matches) < 2 {
			return "", fmt.Errorf("Could not process line: %s", line)
		}

		if !foundFirstIndent {
			foundFirstIndent = true
			globalIndent = matches[1]
		} else {
			if len(globalIndent) > len(matches[1]) {
				globalIndent = matches[1]
			}
		}
	}

	for l, line := range lines {
		if strings.HasPrefix(line, globalIndent) {
			lines[l] = line[len(globalIndent):]
		}
	}

	indentNormalizedTxt := strings.Join(lines, "\n")
	roleRemovedTxt := rstRoleMatcher.ReplaceAllString(indentNormalizedTxt, ":code:`$1`")

	return roleRemovedTxt, nil
}
<|endoftext|>


<commit_before>package main

import (
	"net"
	"os"
	"os/signal"
	"strconv"
	"sync/atomic"
	"time"

	"github.com/pkg/sftp"
	"github.com/spf13/cobra"
	"golang.org/x/sys/unix"
)

/*
#include "config.h"

#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <sched.h>
#include <stdbool.h>
#include <stdio.h>
#include <signal.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>
#include <limits.h>

#include "lxd.h"
#include "memory_utils.h"

void forkfile(void)
{
	int ns_fd = -EBADF, pidfd = -EBADF, rootfs_fd = -EBADF;
	char *listenfd = NULL;
	pid_t pid = 0;

	// Check that we're root.
	if (geteuid() != 0) {
		fprintf(stderr, "Error: forkfile requires root privileges\n");
		_exit(1);
	}

	// Check the first argument.
	listenfd = advance_arg(false);
	if (listenfd == NULL)
		return;

	if (strcmp(listenfd, "--") == 0)
		listenfd = advance_arg(false);

	if (listenfd == NULL || (strcmp(listenfd, "--help") == 0 || strcmp(listenfd, "--version") == 0 || strcmp(listenfd, "-h") == 0))
		return;

	// Get the container rootfs.
	rootfs_fd = atoi(advance_arg(true));

	// Get the container PID.
	pidfd = atoi(advance_arg(true));
	pid = atoi(advance_arg(true));

	if (pid > 0 || pidfd >= 0) {
		ns_fd = pidfd_nsfd(pidfd, pid);
		if (ns_fd < 0) {
			_exit(1);
		}
	}

	// Attach to the container.
	if (ns_fd >= 0) {
		attach_userns_fd(ns_fd);

		if (!change_namespaces(pidfd, ns_fd, CLONE_NEWNS)) {
			error("error: setns");
			_exit(1);
		}
	} else {
		if (fchdir(rootfs_fd) < 0) {
			error("error: fchdir");
			_exit(1);
		}

		if (chroot(".") < 0) {
			error("error: chroot");
			_exit(1);
		}

		if (chdir("/") < 0) {
			error("error: chdir");
			_exit(1);
		}
	}
}
*/
import "C"

type cmdForkfile struct {
	global *cmdGlobal
}

func (c *cmdForkfile) Command() *cobra.Command {
	// Main subcommand
	cmd := &cobra.Command{}
	cmd.Use = "forkfile <listen fd> <rootfs fd> <PIDFd> <PID>"
	cmd.Short = "Perform container file operations"
	cmd.Long = `Description:
  Perform container file operations

  This spawns a daemon inside of the instance's filesystem which can
  then receive command over a simple SFTP API operating on the provided
  listen fd.

  The command can be called with PID and PIDFd set to 0 to just operate on the rootfs fd.
  In such cases, it's the responsibility of the caller to handle any kind of userns shifting.
`
	cmd.Hidden = true
	cmd.Args = cobra.ExactArgs(4)
	cmd.RunE = c.Run

	return cmd
}

func (c *cmdForkfile) Run(cmd *cobra.Command, args []string) error {
	var connections int64

	// Convert the listener FD number.
	listenFD, err := strconv.Atoi(args[0])
	if err != nil {
		return err
	}

	// Setup listener.
	listenerFile := os.NewFile(uintptr(listenFD), "forkfile.sock")
	listener, err := net.FileListener(listenerFile)
	if err != nil {
		return err
	}
	defer func() { _ = listener.Close() }()

	// Convert the rootfs FD number.
	rootfsFD, err := strconv.Atoi(args[1])
	if err != nil {
		return err
	}

	// Automatically shutdown after inactivity.
	go func() {
		for {
			time.Sleep(10 * time.Second)

			if atomic.CompareAndSwapInt64(&connections, 0, -1) {
				// Exit if no more connections.
				os.Exit(0)
			}
		}
	}()

	// Signal handler.
	go func() {
		// Wait for SIGINT.
		sigs := make(chan os.Signal, 1)
		signal.Notify(sigs, unix.SIGINT)
		<-sigs

		// Prevent new connections.
		_ = listener.Close()

		// Wait for connections to be gone and exit.
		for {
			if atomic.CompareAndSwapInt64(&connections, 0, -1) {
				// Exit if no more connections.
				os.Exit(0)
			}

			time.Sleep(time.Second)
		}
	}()

	// Connection handler.
	for {
		// Accept new connection.
		conn, err := listener.Accept()
		if err != nil {
			continue
		}

		go func(conn net.Conn) {
			defer func() { _ = conn.Close() }()

			// Setup connection counter.
			for {
				count := atomic.LoadInt64(&connections)

				// Check if we're currently exiting.
				if count == -1 {
					return
				}

				// Ideally we'd loop here but we can't because go's cmpxchg
				// strangely doesn't return the old value it rather returns a
				// bool preventing such patterns as the one we use here.
				if atomic.CompareAndSwapInt64(&connections, count, count+1) {
					break
				}
			}

			defer func() {
				atomic.AddInt64(&connections, -1)
			}()

			server, err := sftp.NewServer(conn)
			if err != nil {
				return
			}

			_ = server.Serve()

			// Sync the filesystem.
			_ = unix.Syncfs(int(rootfsFD))
		}(conn)
	}
}
<commit_msg>lxd/main_forkfile: Change counting model<commit_after>package main

import (
	"net"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"time"

	"github.com/pkg/sftp"
	"github.com/spf13/cobra"
	"golang.org/x/sys/unix"
)

/*
#include "config.h"

#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <sched.h>
#include <stdbool.h>
#include <stdio.h>
#include <signal.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>
#include <limits.h>

#include "lxd.h"
#include "memory_utils.h"

void forkfile(void)
{
	int ns_fd = -EBADF, pidfd = -EBADF, rootfs_fd = -EBADF;
	char *listenfd = NULL;
	pid_t pid = 0;

	// Check that we're root.
	if (geteuid() != 0) {
		fprintf(stderr, "Error: forkfile requires root privileges\n");
		_exit(1);
	}

	// Check the first argument.
	listenfd = advance_arg(false);
	if (listenfd == NULL)
		return;

	if (strcmp(listenfd, "--") == 0)
		listenfd = advance_arg(false);

	if (listenfd == NULL || (strcmp(listenfd, "--help") == 0 || strcmp(listenfd, "--version") == 0 || strcmp(listenfd, "-h") == 0))
		return;

	// Get the container rootfs.
	rootfs_fd = atoi(advance_arg(true));

	// Get the container PID.
	pidfd = atoi(advance_arg(true));
	pid = atoi(advance_arg(true));

	if (pid > 0 || pidfd >= 0) {
		ns_fd = pidfd_nsfd(pidfd, pid);
		if (ns_fd < 0) {
			_exit(1);
		}
	}

	// Attach to the container.
	if (ns_fd >= 0) {
		attach_userns_fd(ns_fd);

		if (!change_namespaces(pidfd, ns_fd, CLONE_NEWNS)) {
			error("error: setns");
			_exit(1);
		}
	} else {
		if (fchdir(rootfs_fd) < 0) {
			error("error: fchdir");
			_exit(1);
		}

		if (chroot(".") < 0) {
			error("error: chroot");
			_exit(1);
		}

		if (chdir("/") < 0) {
			error("error: chdir");
			_exit(1);
		}
	}
}
*/
import "C"

type cmdForkfile struct {
	global *cmdGlobal
}

func (c *cmdForkfile) Command() *cobra.Command {
	// Main subcommand
	cmd := &cobra.Command{}
	cmd.Use = "forkfile <listen fd> <rootfs fd> <PIDFd> <PID>"
	cmd.Short = "Perform container file operations"
	cmd.Long = `Description:
  Perform container file operations

  This spawns a daemon inside of the instance's filesystem which can
  then receive command over a simple SFTP API operating on the provided
  listen fd.

  The command can be called with PID and PIDFd set to 0 to just operate on the rootfs fd.
  In such cases, it's the responsibility of the caller to handle any kind of userns shifting.
`
	cmd.Hidden = true
	cmd.Args = cobra.ExactArgs(4)
	cmd.RunE = c.Run

	return cmd
}

func (c *cmdForkfile) Run(cmd *cobra.Command, args []string) error {
	var mu sync.RWMutex
	var connections uint64
	var transactions uint64

	// Convert the listener FD number.
	listenFD, err := strconv.Atoi(args[0])
	if err != nil {
		return err
	}

	// Setup listener.
	listenerFile := os.NewFile(uintptr(listenFD), "forkfile.sock")
	listener, err := net.FileListener(listenerFile)
	if err != nil {
		return err
	}
	defer func() { _ = listener.Close() }()

	// Convert the rootfs FD number.
	rootfsFD, err := strconv.Atoi(args[1])
	if err != nil {
		return err
	}

	// Automatically shutdown after inactivity.
	go func() {
		for {
			time.Sleep(10 * time.Second)

			// Check for active connections.
			mu.RLock()
			if connections > 0 {
				mu.RUnlock()
				continue
			}

			// Look for recent activity
			oldCount := transactions
			mu.RUnlock()

			time.Sleep(5 * time.Second)

			mu.RLock()
			if oldCount == transactions {
				mu.RUnlock()

				// Daemon has been inactive for 10s, exit.
				os.Exit(0)
			}
			mu.RUnlock()
		}
	}()

	// Signal handler.
	go func() {
		// Wait for SIGINT.
		sigs := make(chan os.Signal, 1)
		signal.Notify(sigs, unix.SIGINT)
		<-sigs

		// Prevent new connections.
		_ = listener.Close()

		// Wait for connections to be gone and exit.
		for {
			mu.RLock()
			if connections == 0 {
				mu.RUnlock()
				break
			}
			mu.RUnlock()

			time.Sleep(time.Second)
		}

		os.Exit(0)
	}()

	// Connection handler.
	for {
		// Accept new connection.
		conn, err := listener.Accept()
		if err != nil {
			continue
		}

		go func(conn net.Conn) {
			defer func() {
				_ = conn.Close()

				mu.Lock()
				connections -= 1
				mu.Unlock()
			}()

			// Increase counters.
			mu.Lock()
			transactions += 1
			connections += 1
			mu.Unlock()

			// Spawn the server.
			server, err := sftp.NewServer(conn)
			if err != nil {
				return
			}

			_ = server.Serve()

			// Sync the filesystem.
			_ = unix.Syncfs(int(rootfsFD))
		}(conn)
	}
}
<|endoftext|>


<commit_before>package codewriter

import (
	"bufio"
	"fmt"
	"io"
	"strconv"
	"strings"
	"sync"
)

// baseLabel is a base name of labels.
const baseLabel = "LABEL"

// Mneumonic is a mneumonic of an instruction.
type Mneumonic string

// Mneumonics.
const (
	Add  Mneumonic = "add"
	Sub            = "sub"
	Neg            = "neg"
	Eq             = "eq"
	Gt             = "gt"
	Lt             = "lt"
	And            = "and"
	Or             = "or"
	Not            = "not"
	Push           = "push"
	Pop            = "pop"
)

// Segment is a memory segment.
type Segment string

// Memory segments.
const (
	Argument Segment = "argument"
	Local            = "local"
	Static           = "static"
	Constant         = "constant"
	This             = "this"
	That             = "that"
	Pointer          = "pointer"
	Temp             = "temp"
)

// CodeWriter converts VM commands to Hack assembly codes and write them out to a destination.
type CodeWriter struct {
	err      error
	dest     io.Writer
	buf      *bufio.Writer
	filename string

	mu  sync.Mutex
	cnt int
}

// New creates a new CodeWriter that writes converted codes to dest.
func New(dest io.Writer) *CodeWriter {
	return &CodeWriter{
		dest: dest,
		buf:  bufio.NewWriter(dest),
	}
}

// SetFileName sets an input VM file name to filename.
func (cw *CodeWriter) SetFileName(filename string) {
	cw.filename = filename
}

// WriteArithmetic converts the given arithmetic command to assembly code and writes it out.
func (cw *CodeWriter) WriteArithmetic(cmd string) error {
	switch cmd {
	case "neg", "not":
		cw.unary(cmd)
	case "add", "sub", "and", "or":
		cw.binary(cmd)
	case "eq", "gt", "lt":
		cw.compare(cmd)
	default:
		cw.err = fmt.Errorf("unknown command: %s", cmd)
	}

	if cw.err != nil {
		return fmt.Errorf("error writing code: %s", cw.err.Error())
	}
	return nil
}

// Close flushes bufferred data to the destination and closes it.
// Note that no data is written to the destination until Close is called.
func (cw *CodeWriter) Close() error {
	defer func() {
		if cl, ok := cw.dest.(io.Closer); ok {
			_ = cl.Close()
		}
	}()

	if e := cw.buf.Flush(); e != nil {
		return fmt.Errorf("error flushing bufferred data: %s", e)
	}
	return nil
}

// unary writes a unary operation for a value at the top of the stack.
// cmd must be one of the following:
//   - "neg"
//   - "not"
func (cw *CodeWriter) unary(cmd string) {
	var op string
	switch cmd {
	case "neg":
		op = "-"
	case "not":
		op = "!"
	}

	cw.decrSP()
	cw.ccmd("M", op+"M")
	cw.incrSP()
}

// binary writes a binary operation for two values at the top of the stack.
// cmd must be one of the following:
//   - "add"
//   - "sub"
//   - "and"
//   - "or"
func (cw *CodeWriter) binary(cmd string) {
	var op string
	switch cmd {
	case "add":
		op = "M+D"
	case "sub":
		op = "M-D"
	case "and":
		op = "D&M"
	case "or":
		op = "D|M"
	}

	cw.popStack()
	cw.decrSP()
	cw.ccmd("M", op)
	cw.incrSP()
}

// compare writes a comparison operation for two values at the top of the stack.
// cmd must be one of the following:
//   - "eq"
//   - "gt"
//   - "lt"
func (cw *CodeWriter) compare(cmd string) {
	// JEQ, JGT, JLT
	op := "J" + strings.ToUpper(cmd)
	label1, label2 := cw.label(), cw.label()

	cw.popStack()
	cw.decrSP()
	cw.ccmd("D", "M-D")
	cw.acmd(label1)
	cw.ccmdj("", "D", op)
	cw.loadToSP("0")
	cw.acmd(label2)
	cw.ccmdj("", "0", "JMP")
	cw.lcmd(label1)
	cw.loadToSP("-1")
	cw.lcmd(label2)
	cw.incrSP()
}

// label returns a label.
func (cw *CodeWriter) label() string {
	defer cw.countUp()
	return baseLabel + strconv.Itoa(cw.cnt)
}

// countUp counts up an internal counter.
func (cw *CodeWriter) countUp() {
	cw.mu.Lock()
	defer cw.mu.Unlock()
	cw.cnt++
}

// pushStack pushes v to the top of the stack. Internally,
// it assgins v to *SP and increments SP.
// If an error occurs and cw.err is nil, it is set at cw.err.
func (cw *CodeWriter) pushStack(v string) {
	cw.loadToSP(v)
	cw.incrSP()
}

// loadToSP loads v to *SP.
func (cw *CodeWriter) loadToSP(v string) {
	cw.acmd("SP")
	cw.ccmd("A", "M")
	cw.ccmd("M", v)
}

// popStack pops a value at the top of the stack. Internally,
// it decrements SP and assigns a value pointed by SP to D.
// If an error occurs and cw.err is nil, it is set at cw.err.
func (cw *CodeWriter) popStack() {
	cw.decrSP()
	cw.ccmd("D", "M")
}

// incrSP increments SP and sets the current address to it.
// If an error occurs and cw.err is nil, it is set at cw.err.
func (cw *CodeWriter) incrSP() {
	cw.sp("+")
}

// decrSP increments SP and sets the current address to it.
// If an error occurs and cw.err is nil, it is set at cw.err.
func (cw *CodeWriter) decrSP() {
	cw.sp("-")
}

// sp controls the position of SP and sets the current address to it.
// op must be one of the following:
//   "+": SP++
//   "-": SP--
func (cw *CodeWriter) sp(op string) {
	cw.acmd("SP")
	cw.ccmd("AM", "M"+op+"1")
}

// acmd writes @ command. If an error occurs and cw.err is nil, it is set at cw.err.
func (cw *CodeWriter) acmd(addr string) {
	if _, e := cw.buf.WriteString("@" + addr + "\n"); cw.err == nil {
		cw.err = e
	}
}

// ccmd writes C command with no jump. If an error occurs, it is set at cw.err.
func (cw *CodeWriter) ccmd(dest, comp string) {
	cw.ccmdj(dest, comp, "")
}

// ccmdj writes C command with jump. If an error occurs, it is set at cw.err.
func (cw *CodeWriter) ccmdj(dest, comp, jump string) {
	// allocate a slice whose length is len(dest=comp;jump\n)
	opc := make([]byte, 0, len(dest)+1+len(comp)+1+len(jump)+1)

	// append `dest=`
	if dest != "" {
		opc = append(append(opc, dest...), '=')
	}

	// append comp
	opc = append(opc, comp...)

	// append `;jump`
	if jump != "" {
		opc = append(append(opc, ';'), jump...)
	}

	// append \n
	opc = append(opc, '\n')

	if _, e := cw.buf.Write(opc); cw.err == nil {
		cw.err = e
	}
}

// lcmd writes label command. If an error occurs, it is set at cw.err.
func (cw *CodeWriter) lcmd(label string) {
	if _, e := cw.buf.WriteString("(" + label + ")\n"); cw.err == nil {
		cw.err = e
	}
}
<commit_msg>codewriter: return at first if CodeWriter.err is not nil<commit_after>package codewriter

import (
	"bufio"
	"fmt"
	"io"
	"strconv"
	"strings"
	"sync"
)

// baseLabel is a base name of labels.
const baseLabel = "LABEL"

// Mneumonic is a mneumonic of an instruction.
type Mneumonic string

// Mneumonics.
const (
	Add  Mneumonic = "add"
	Sub            = "sub"
	Neg            = "neg"
	Eq             = "eq"
	Gt             = "gt"
	Lt             = "lt"
	And            = "and"
	Or             = "or"
	Not            = "not"
	Push           = "push"
	Pop            = "pop"
)

// Segment is a memory segment.
type Segment string

// Memory segments.
const (
	Argument Segment = "argument"
	Local            = "local"
	Static           = "static"
	Constant         = "constant"
	This             = "this"
	That             = "that"
	Pointer          = "pointer"
	Temp             = "temp"
)

// CodeWriter converts VM commands to Hack assembly codes and write them out to a destination.
type CodeWriter struct {
	err      error
	dest     io.Writer
	buf      *bufio.Writer
	filename string

	mu  sync.Mutex
	cnt int
}

// New creates a new CodeWriter that writes converted codes to dest.
func New(dest io.Writer) *CodeWriter {
	return &CodeWriter{
		dest: dest,
		buf:  bufio.NewWriter(dest),
	}
}

// SetFileName sets an input VM file name to filename.
func (cw *CodeWriter) SetFileName(filename string) {
	cw.filename = filename
}

// WriteArithmetic converts the given arithmetic command to assembly code and writes it out.
func (cw *CodeWriter) WriteArithmetic(cmd string) error {
	switch cmd {
	case "neg", "not":
		cw.unary(cmd)
	case "add", "sub", "and", "or":
		cw.binary(cmd)
	case "eq", "gt", "lt":
		cw.compare(cmd)
	default:
		cw.err = fmt.Errorf("unknown command: %s", cmd)
	}

	if cw.err != nil {
		return fmt.Errorf("error writing code: %s", cw.err.Error())
	}
	return nil
}

// Close flushes bufferred data to the destination and closes it.
// Note that no data is written to the destination until Close is called.
func (cw *CodeWriter) Close() error {
	defer func() {
		if cl, ok := cw.dest.(io.Closer); ok {
			_ = cl.Close()
		}
	}()

	if e := cw.buf.Flush(); e != nil {
		return fmt.Errorf("error flushing bufferred data: %s", e)
	}
	return nil
}

// unary writes a unary operation for a value at the top of the stack.
// cmd must be one of the following:
//   - "neg"
//   - "not"
func (cw *CodeWriter) unary(cmd string) {
	var op string
	switch cmd {
	case "neg":
		op = "-"
	case "not":
		op = "!"
	}

	cw.decrSP()
	cw.ccmd("M", op+"M")
	cw.incrSP()
}

// binary writes a binary operation for two values at the top of the stack.
// cmd must be one of the following:
//   - "add"
//   - "sub"
//   - "and"
//   - "or"
func (cw *CodeWriter) binary(cmd string) {
	var op string
	switch cmd {
	case "add":
		op = "M+D"
	case "sub":
		op = "M-D"
	case "and":
		op = "D&M"
	case "or":
		op = "D|M"
	}

	cw.popStack()
	cw.decrSP()
	cw.ccmd("M", op)
	cw.incrSP()
}

// compare writes a comparison operation for two values at the top of the stack.
// cmd must be one of the following:
//   - "eq"
//   - "gt"
//   - "lt"
func (cw *CodeWriter) compare(cmd string) {
	// JEQ, JGT, JLT
	op := "J" + strings.ToUpper(cmd)
	label1, label2 := cw.label(), cw.label()

	cw.popStack()
	cw.decrSP()
	cw.ccmd("D", "M-D")
	cw.acmd(label1)
	cw.ccmdj("", "D", op)
	cw.loadToSP("0")
	cw.acmd(label2)
	cw.ccmdj("", "0", "JMP")
	cw.lcmd(label1)
	cw.loadToSP("-1")
	cw.lcmd(label2)
	cw.incrSP()
}

// label returns a label.
func (cw *CodeWriter) label() string {
	defer cw.countUp()
	return baseLabel + strconv.Itoa(cw.cnt)
}

// countUp counts up an internal counter.
func (cw *CodeWriter) countUp() {
	cw.mu.Lock()
	defer cw.mu.Unlock()
	cw.cnt++
}

// pushStack pushes v to the top of the stack. Internally,
// it assgins v to *SP and increments SP.
// If an error occurs and cw.err is nil, it is set at cw.err.
func (cw *CodeWriter) pushStack(v string) {
	cw.loadToSP(v)
	cw.incrSP()
}

// loadToSP loads v to *SP.
func (cw *CodeWriter) loadToSP(v string) {
	cw.acmd("SP")
	cw.ccmd("A", "M")
	cw.ccmd("M", v)
}

// popStack pops a value at the top of the stack. Internally,
// it decrements SP and assigns a value pointed by SP to D.
// If an error occurs and cw.err is nil, it is set at cw.err.
func (cw *CodeWriter) popStack() {
	cw.decrSP()
	cw.ccmd("D", "M")
}

// incrSP increments SP and sets the current address to it.
// If an error occurs and cw.err is nil, it is set at cw.err.
func (cw *CodeWriter) incrSP() {
	cw.sp("+")
}

// decrSP increments SP and sets the current address to it.
// If an error occurs and cw.err is nil, it is set at cw.err.
func (cw *CodeWriter) decrSP() {
	cw.sp("-")
}

// sp controls the position of SP and sets the current address to it.
// op must be one of the following:
//   "+": SP++
//   "-": SP--
func (cw *CodeWriter) sp(op string) {
	cw.acmd("SP")
	cw.ccmd("AM", "M"+op+"1")
}

// acmd writes @ command. If an error occurs and cw.err is nil, it is set at cw.err.
func (cw *CodeWriter) acmd(addr string) {
	if cw.err != nil {
		return
	}

	_, cw.err = cw.buf.WriteString("@" + addr + "\n")
}

// ccmd writes C command with no jump. If an error occurs, it is set at cw.err.
func (cw *CodeWriter) ccmd(dest, comp string) {
	cw.ccmdj(dest, comp, "")
}

// ccmdj writes C command with jump. If an error occurs, it is set at cw.err.
func (cw *CodeWriter) ccmdj(dest, comp, jump string) {
	if cw.err != nil {
		return
	}

	// allocate a slice whose length is len(dest=comp;jump\n)
	opc := make([]byte, 0, len(dest)+1+len(comp)+1+len(jump)+1)

	// append `dest=`
	if dest != "" {
		opc = append(append(opc, dest...), '=')
	}

	// append comp
	opc = append(opc, comp...)

	// append `;jump`
	if jump != "" {
		opc = append(append(opc, ';'), jump...)
	}

	// append \n
	opc = append(opc, '\n')

	_, cw.err = cw.buf.Write(opc)
}

// lcmd writes label command. If an error occurs, it is set at cw.err.
func (cw *CodeWriter) lcmd(label string) {
	if cw.err != nil {
		return
	}

	_, cw.err = cw.buf.WriteString("(" + label + ")\n")
}
<|endoftext|>


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

// A program that makes random reads within a large object on GCS, reporting
// throughput and latency information.
package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"os/signal"
	"sync"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"

	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/syncutil"
)

var fBucket = flag.String("bucket", "", "The GCS bucket from which to read.")
var fObject = flag.String("object", "", "The object within which to read.")
var fSize = flag.Uint64("size", 1<<21, "The size of each read in bytes.")
var fWorkers = flag.Int("workers", 1, "The number of workers to run.")

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type result struct {
	BytesRead int

	// Time taken from making the HTTP request to receiving the first byte of the
	// response body.
	FirstByteLatency time.Duration

	// Time taken from making the HTTP request to receiving the last byte of the
	// response body.
	FullBodyDuration time.Duration
}

func readOnce(
	ctx context.Context,
	o *gcs.Object,
	bucket gcs.Bucket) (r result, err error) {
	// Is the object large enough?
	if o.Size < *fSize {
		err = fmt.Errorf(
			"Object of size %d not large enough for read size %d",
			o.Size,
			*fSize)

		return
	}

	// Set up an appropriate request.
	req := &gcs.ReadObjectRequest{
		Name:       o.Name,
		Generation: o.Generation,
		Range:      &gcs.ByteRange{},
	}

	req.Range.Start = uint64(rand.Int63n(int64(o.Size - *fSize)))
	req.Range.Limit = req.Range.Start + *fSize

	// Create the reader.
	start := time.Now()
	rc, err := bucket.NewReader(ctx, req)
	if err != nil {
		err = fmt.Errorf("NewReader: %v", err)
		return
	}

	defer func() {
		closeErr := rc.Close()
		if err == nil && closeErr != nil {
			err = fmt.Errorf("Close: %v", closeErr)
		}
	}()

	// Measure the time to first byte.
	_, err = rc.Read([]byte{0})
	if err != nil {
		err = fmt.Errorf("Read: %v", err)
		return
	}

	r.FirstByteLatency = time.Since(start)

	// And the time to read everything.
	n, err := io.Copy(ioutil.Discard, rc)
	if err != nil {
		err = fmt.Errorf("Copy: %v", err)
		return
	}

	r.FullBodyDuration = time.Since(start)
	r.BytesRead = int(n + 1)

	return
}

// Make random reads within the given object, writing results to the supplied
// channel. Stop when the stop channel is closed.
func makeReads(
	ctx context.Context,
	o *gcs.Object,
	bucket gcs.Bucket,
	results chan<- result,
	stop <-chan struct{}) (err error) {
	// Is the object large enough?
	if o.Size < *fSize {
		err = fmt.Errorf(
			"Object of size %d not large enough for read size %d",
			o.Size,
			*fSize)

		return
	}

	for {
		// Stop?
		if _, ok := <-stop; ok {
			return
		}

		var r result
		r, err = readOnce(ctx, o, bucket)
	}
}

func getBucket(ctx context.Context) (b gcs.Bucket, err error) {
	if *fBucket == "" {
		err = errors.New("You must set --bucket.")
		return
	}

	// Set up the token source.
	const scope = gcs.Scope_ReadOnly
	tokenSrc, err := google.DefaultTokenSource(context.Background(), scope)
	if err != nil {
		err = fmt.Errorf("DefaultTokenSource: %v", err)
		return
	}

	// Use that to create a GCS connection.
	cfg := &gcs.ConnConfig{
		TokenSource: tokenSrc,
	}

	conn, err := gcs.NewConn(cfg)
	if err != nil {
		err = fmt.Errorf("NewConn: %v", err)
		return
	}

	// Open the bucket.
	b, err = conn.OpenBucket(ctx, *fBucket)
	if err != nil {
		err = fmt.Errorf("OpenBucket: %v", err)
		return
	}

	return
}

func getObject(
	ctx context.Context,
	bucket gcs.Bucket) (o *gcs.Object, err error) {
	if *fObject == "" {
		err = errors.New("You must set --object.")
		return
	}

	o, err = bucket.StatObject(
		ctx,
		&gcs.StatObjectRequest{Name: *fObject})

	return
}

// Run workers until SIGINT is received. Return a slice of results.
func runWorkers(
	ctx context.Context,
	o *gcs.Object,
	bucket gcs.Bucket) (results []result, err error) {
	b := syncutil.NewBundle(ctx)

	// Set up a channel that is closed upon SIGINT.
	stop := make(chan struct{})
	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, os.Interrupt)

		<-c
		log.Printf("SIGINT received. Stopping soon...")
		close(stop)
	}()

	// Start several workers making random reads.
	var wg sync.WaitGroup
	resultChan := make(chan result)
	for i := 0; i < *fWorkers; i++ {
		wg.Add(1)
		b.Add(func(ctx context.Context) (err error) {
			defer wg.Done()
			err = makeReads(ctx, o, bucket, resultChan, stop)
			if err != nil {
				err = fmt.Errorf("makeReads: %v", err)
				return
			}

			return
		})
	}

	go func() {
		wg.Wait()
		close(resultChan)
	}()

	// Accumulate results.
	b.Add(func(ctx context.Context) (err error) {
		for r := range resultChan {
			results = append(results, r)
		}

		return
	})

	err = b.Join()
	return
}

func describeResults(results []result)

////////////////////////////////////////////////////////////////////////
// Main
////////////////////////////////////////////////////////////////////////

func run(ctx context.Context) (err error) {
	// Open the bucket.
	bucket, err := getBucket(ctx)
	if err != nil {
		err = fmt.Errorf("getBucket: %v", err)
		return
	}

	// Stat the object.
	o, err := getObject(ctx, bucket)
	if err != nil {
		err = fmt.Errorf("getObject: %v", err)
		return
	}

	// Make reads.
	results, err := runWorkers(ctx, o, bucket)
	if err != nil {
		err = fmt.Errorf("runWorkers: %v", err)
		return
	}

	// Print information about the results.
	describeResults(results)

	return
}

func main() {
	log.SetFlags(log.Lmicroseconds)

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

// A program that makes random reads within a large object on GCS, reporting
// throughput and latency information.
package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"os/signal"
	"sync"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"

	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/syncutil"
)

var fBucket = flag.String("bucket", "", "The GCS bucket from which to read.")
var fObject = flag.String("object", "", "The object within which to read.")
var fSize = flag.Uint64("size", 1<<21, "The size of each read in bytes.")
var fWorkers = flag.Int("workers", 1, "The number of workers to run.")

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type result struct {
	BytesRead int

	// Time taken from making the HTTP request to receiving the first byte of the
	// response body.
	FirstByteLatency time.Duration

	// Time taken from making the HTTP request to receiving the last byte of the
	// response body.
	FullBodyDuration time.Duration
}

func readOnce(
	ctx context.Context,
	o *gcs.Object,
	bucket gcs.Bucket) (r result, err error) {
	// Is the object large enough?
	if o.Size < *fSize {
		err = fmt.Errorf(
			"Object of size %d not large enough for read size %d",
			o.Size,
			*fSize)

		return
	}

	// Set up an appropriate request.
	req := &gcs.ReadObjectRequest{
		Name:       o.Name,
		Generation: o.Generation,
		Range:      &gcs.ByteRange{},
	}

	req.Range.Start = uint64(rand.Int63n(int64(o.Size - *fSize)))
	req.Range.Limit = req.Range.Start + *fSize

	// Create the reader.
	start := time.Now()
	rc, err := bucket.NewReader(ctx, req)
	if err != nil {
		err = fmt.Errorf("NewReader: %v", err)
		return
	}

	defer func() {
		closeErr := rc.Close()
		if err == nil && closeErr != nil {
			err = fmt.Errorf("Close: %v", closeErr)
		}
	}()

	// Measure the time to first byte.
	_, err = rc.Read([]byte{0})
	if err != nil {
		err = fmt.Errorf("Read: %v", err)
		return
	}

	r.FirstByteLatency = time.Since(start)

	// And the time to read everything.
	n, err := io.Copy(ioutil.Discard, rc)
	if err != nil {
		err = fmt.Errorf("Copy: %v", err)
		return
	}

	r.FullBodyDuration = time.Since(start)
	r.BytesRead = int(n + 1)

	return
}

// Make random reads within the given object, writing results to the supplied
// channel. Stop when the stop channel is closed.
func makeReads(
	ctx context.Context,
	o *gcs.Object,
	bucket gcs.Bucket,
	results chan<- result,
	stop <-chan struct{}) (err error) {
	// Is the object large enough?
	if o.Size < *fSize {
		err = fmt.Errorf(
			"Object of size %d not large enough for read size %d",
			o.Size,
			*fSize)

		return
	}

	for {
		// Stop?
		if _, ok := <-stop; ok {
			return
		}

		// Make a read.
		var r result
		r, err = readOnce(ctx, o, bucket)
		if err != nil {
			err = fmt.Errorf("readOnce %v", err)
			return
		}

		// Write out the result.
		select {
		case <-ctx.Done():
			err = ctx.Err()
			return

		case results <- r:
		}
	}
}

func getBucket(ctx context.Context) (b gcs.Bucket, err error) {
	if *fBucket == "" {
		err = errors.New("You must set --bucket.")
		return
	}

	// Set up the token source.
	const scope = gcs.Scope_ReadOnly
	tokenSrc, err := google.DefaultTokenSource(context.Background(), scope)
	if err != nil {
		err = fmt.Errorf("DefaultTokenSource: %v", err)
		return
	}

	// Use that to create a GCS connection.
	cfg := &gcs.ConnConfig{
		TokenSource: tokenSrc,
	}

	conn, err := gcs.NewConn(cfg)
	if err != nil {
		err = fmt.Errorf("NewConn: %v", err)
		return
	}

	// Open the bucket.
	b, err = conn.OpenBucket(ctx, *fBucket)
	if err != nil {
		err = fmt.Errorf("OpenBucket: %v", err)
		return
	}

	return
}

func getObject(
	ctx context.Context,
	bucket gcs.Bucket) (o *gcs.Object, err error) {
	if *fObject == "" {
		err = errors.New("You must set --object.")
		return
	}

	o, err = bucket.StatObject(
		ctx,
		&gcs.StatObjectRequest{Name: *fObject})

	return
}

// Run workers until SIGINT is received. Return a slice of results.
func runWorkers(
	ctx context.Context,
	o *gcs.Object,
	bucket gcs.Bucket) (results []result, err error) {
	b := syncutil.NewBundle(ctx)

	// Set up a channel that is closed upon SIGINT.
	stop := make(chan struct{})
	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, os.Interrupt)

		<-c
		log.Printf("SIGINT received. Stopping soon...")
		close(stop)
	}()

	// Start several workers making random reads.
	var wg sync.WaitGroup
	resultChan := make(chan result)
	for i := 0; i < *fWorkers; i++ {
		wg.Add(1)
		b.Add(func(ctx context.Context) (err error) {
			defer wg.Done()
			err = makeReads(ctx, o, bucket, resultChan, stop)
			if err != nil {
				err = fmt.Errorf("makeReads: %v", err)
				return
			}

			return
		})
	}

	go func() {
		wg.Wait()
		close(resultChan)
	}()

	// Accumulate results.
	b.Add(func(ctx context.Context) (err error) {
		for r := range resultChan {
			results = append(results, r)
		}

		return
	})

	err = b.Join()
	return
}

func describeResults(results []result)

////////////////////////////////////////////////////////////////////////
// Main
////////////////////////////////////////////////////////////////////////

func run(ctx context.Context) (err error) {
	// Open the bucket.
	bucket, err := getBucket(ctx)
	if err != nil {
		err = fmt.Errorf("getBucket: %v", err)
		return
	}

	// Stat the object.
	o, err := getObject(ctx, bucket)
	if err != nil {
		err = fmt.Errorf("getObject: %v", err)
		return
	}

	// Make reads.
	results, err := runWorkers(ctx, o, bucket)
	if err != nil {
		err = fmt.Errorf("runWorkers: %v", err)
		return
	}

	// Print information about the results.
	describeResults(results)

	return
}

func main() {
	log.SetFlags(log.Lmicroseconds)

	err := run(context.Background())
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package mailjet

import (
	"net/http"
)

// MailjetClient bundles data needed by a large number
// of methods in order to interact with the Mailjet API.
type MailjetClient struct {
	apiKeyPublic  string
	apiKeyPrivate string
	client        *http.Client
}

// MailjetRequest bundles data needed to build the URL.
type MailjetRequest struct {
	Resource string
	ID       int64
	AltID    string
	Action   string
	ActionID int64
}

// MailjetDataRequest bundles data needed to build the DATA URL.
type MailjetDataRequest struct {
	SourceType   string
	SourceTypeID int64
	DataType     string
	MimeType     string
	DataTypeID   int64
	LastID       bool
}

// FullMailjetRequest is the same as a MailjetRequest but with a payload.
type FullMailjetRequest struct {
	Info    *MailjetRequest
	Payload interface{}
}

// FullMailjetDataRequest is the same as a MailjetDataRequest but with a payload.
type FullMailjetDataRequest struct {
	Info    *MailjetDataRequest
	Payload interface{}
}

type MailjetOptions func(*http.Request)

// MailjetResult is the JSON result sent by the API.
type MailjetResult struct {
	Count int
	Data  interface{}
	Total int
}

// MailjetError is the error returned by the API.
type MailjetError struct {
	ErrorInfo    string
	ErrorMessage string
	StatusCode   int
}

//
// Send API structures
//

type MailjetSendMail struct {
	FromEmail             string
	FromName              string
	Sender                string
	Recipients            []MailjetRecipient
	To                    []string
	Cc                    []string
	Bcc                   []string
	Subject               string
	TextPart              string `json:"Text-part"`
	HtmlPart              string `json:"Html-part"`
	Attachments           []MailjetAttachment
	InlineAttachments     []MailjetAttachment `json:"Inline_attachments"`
	MjPrio                int                 `json:"Mj-prio"`
	MjCampaign            string              `json:"Mj-campaign"`
	MjDeduplicateCampaign bool                `json:"Mj-deduplicatecampaign"`
	MjCustomID            string              `json:"Mj-CustomID"`
	MjEventPayload        string              `json:"Mj-EventPayLoad"`
	MjTemplateID	      string		  `json:"Mj-Template-ID"`
	Headers               map[string]string
	Vars                  interface{}
	Messages              []MailjetSendMail
}

type MailjetRecipient struct {
	Email string
	Name  string
	Vars  interface{}
}

type MailjetAttachment struct {
	ContentType string `json:"Content-Type"`
	Content     string
	Filename    string
}

type MailjetSent struct {
	Email     string
	MessageID int
}

type MailjetSentResult struct {
	Sent []MailjetSent
}
<commit_msg>haeders temp fix<commit_after>package mailjet

import (
	"net/http"
)

// MailjetClient bundles data needed by a large number
// of methods in order to interact with the Mailjet API.
type MailjetClient struct {
	apiKeyPublic  string
	apiKeyPrivate string
	client        *http.Client
}

// MailjetRequest bundles data needed to build the URL.
type MailjetRequest struct {
	Resource string
	ID       int64
	AltID    string
	Action   string
	ActionID int64
}

// MailjetDataRequest bundles data needed to build the DATA URL.
type MailjetDataRequest struct {
	SourceType   string
	SourceTypeID int64
	DataType     string
	MimeType     string
	DataTypeID   int64
	LastID       bool
}

// FullMailjetRequest is the same as a MailjetRequest but with a payload.
type FullMailjetRequest struct {
	Info    *MailjetRequest
	Payload interface{}
}

// FullMailjetDataRequest is the same as a MailjetDataRequest but with a payload.
type FullMailjetDataRequest struct {
	Info    *MailjetDataRequest
	Payload interface{}
}

type MailjetOptions func(*http.Request)

// MailjetResult is the JSON result sent by the API.
type MailjetResult struct {
	Count int
	Data  interface{}
	Total int
}

// MailjetError is the error returned by the API.
type MailjetError struct {
	ErrorInfo    string
	ErrorMessage string
	StatusCode   int
}

//
// Send API structures
//

type MailjetSendMail struct {
	FromEmail             string
	FromName              string
	Sender                string
	Recipients            []MailjetRecipient
	To                    []string
	Cc                    []string
	Bcc                   []string
	Subject               string
	TextPart              string `json:"Text-part"`
	HtmlPart              string `json:"Html-part"`
	Attachments           []MailjetAttachment
	InlineAttachments     []MailjetAttachment `json:"Inline_attachments"`
	MjPrio                int                 `json:"Mj-prio"`
	MjCampaign            string              `json:"Mj-campaign"`
	MjDeduplicateCampaign bool                `json:"Mj-deduplicatecampaign"`
	MjCustomID            string              `json:"Mj-CustomID"`
	MjEventPayload        string              `json:"Mj-EventPayLoad"`
	MjTemplateID          string              `json:"Mj-Template-ID"`
	Headers               interface{}         `json:",omitempty"`
	Vars                  interface{}
	Messages              []MailjetSendMail
}

type MailjetRecipient struct {
	Email string
	Name  string
	Vars  interface{}
}

type MailjetAttachment struct {
	ContentType string `json:"Content-Type"`
	Content     string
	Filename    string
}

type MailjetSent struct {
	Email     string
	MessageID int
}

type MailjetSentResult struct {
	Sent []MailjetSent
}
<|endoftext|>


<commit_before>package dropp

import (
	"encoding/json"
	"log"
	"net/http"
	"time"

	gq "github.com/PuerkitoBio/goquery"
)

// Snapshot contains a snapshot of the current status of an item.
type Snapshot struct {
	Availability AvaComp
	OnEbay       bool
	Price        float64
	CreatedAt    time.Time
}

// SnapshotDiff Is created if there is a difference between the current and the
// previous snaphot.
type SnapshotDiff struct {
	ItemName       string
	ItemURL        string
	PreviousAva    AvaComp
	PreviousStatus bool
	PreviousPrice  float64
	CurrentAva     AvaComp
	CurrentStatus  bool
	CurrentPrice   float64
}

// EbayPrice contains the price and currency as fetched from the Ebay service
type EbayPrice struct {
	CurrencyID string `json:"_currencyID"`
	Value      string `json:"value"`
}

// EbayItem is the full item generated by the Ebay service
type EbayItem struct {
	CurrentPrice EbayPrice `json:"current_price"`
	ID           string    `json:"id"`
	ItemName     string    `json:"name"`
	StockCount   string    `json:"quantity"`
	SoldCount    string    `json:"quantity_sold"`
	Status       string    `json:"status"`
	ItemPageURL  string    `json:"url"`
}

func (snap *Snapshot) getBGAva(response *http.Response) {
	// Scrape the page and get availability
	log.Print("Scraping BG page to retrieve availability")
	doc, err := gq.NewDocumentFromResponse(response)
	if err != nil {
		panic(err.Error())
	}
	ava := doc.Find(".status").Text()
	snap.Availability = NewAva(ava)
	return
}

func (snap *Snapshot) getEbayStatus(response *http.Response) {
	ebayItem := EbayItem{}

	err := json.NewDecoder(response.Body).Decode(&ebayItem)
	if err != nil {
		log.Printf("Error while converting the Ebay JSON %s", err)
		return
	}

	log.Printf("The status of item %s is %s", ebayItem.ID, ebayItem.Status)

	if ebayItem.Status == "Active" {
		snap.OnEbay = true
	} else {
		snap.OnEbay = false
	}
}
<commit_msg>[DPP-92] Changed availability type in snapshot types<commit_after>package dropp

import (
	"encoding/json"
	"log"
	"net/http"
	"time"

	gq "github.com/PuerkitoBio/goquery"
)

// Snapshot contains a snapshot of the current status of an item.
type Snapshot struct {
	Availability string
	OnEbay       bool
	Price        float64
	CreatedAt    time.Time
}

// SnapshotDiff Is created if there is a difference between the current and the
// previous snaphot.
type SnapshotDiff struct {
	ItemName       string
	ItemURL        string
	PreviousAva    string
	PreviousStatus bool
	PreviousPrice  float64
	CurrentAva     string
	CurrentStatus  bool
	CurrentPrice   float64
}

// EbayPrice contains the price and currency as fetched from the Ebay service
type EbayPrice struct {
	CurrencyID string `json:"_currencyID"`
	Value      string `json:"value"`
}

// EbayItem is the full item generated by the Ebay service
type EbayItem struct {
	CurrentPrice EbayPrice `json:"current_price"`
	ID           string    `json:"id"`
	ItemName     string    `json:"name"`
	StockCount   string    `json:"quantity"`
	SoldCount    string    `json:"quantity_sold"`
	Status       string    `json:"status"`
	ItemPageURL  string    `json:"url"`
}

func (snap *Snapshot) getBGAva(response *http.Response) {
	// Scrape the page and get availability
	log.Print("Scraping BG page to retrieve availability")
	doc, err := gq.NewDocumentFromResponse(response)
	if err != nil {
		panic(err.Error())
	}
	ava := doc.Find(".status").Text()
	snap.Availability = ava
	return
}

func (snap *Snapshot) getEbayStatus(response *http.Response) {
	ebayItem := EbayItem{}

	err := json.NewDecoder(response.Body).Decode(&ebayItem)
	if err != nil {
		log.Printf("Error while converting the Ebay JSON %s", err)
		return
	}

	log.Printf("The status of item %s is %s", ebayItem.ID, ebayItem.Status)

	if ebayItem.Status == "Active" {
		snap.OnEbay = true
	} else {
		snap.OnEbay = false
	}
}
<|endoftext|>


<commit_before>// Copyright © 2017 Ricardo Aravena <raravena@branch.io>
//
// 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 exec

import (
	"bytes"
	"fmt"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"io/ioutil"
	"net"
	"os"
	"time"
)

type Options struct {
	machines []string
	port     string
	user     string
	cmd      string
	key      string
	useAgent bool
}

type executeResult struct {
	result string
	err    error
}

func User(u string) func(*Options) {
	return func(e *Options) {
		e.user = u
	}
}

func Port(p string) func(*Options) {
	return func(e *Options) {
		e.port = p
	}
}

func Cmd(c string) func(*Options) {
	return func(e *Options) {
		e.cmd = c
	}
}

func Machines(m []string) func(*Options) {
	return func(e *Options) {
		e.machines = m
	}
}

func Key(k string) func(*Options) {
	return func(e *Options) {
		e.key = k
	}
}

func UseAgent(u bool) func(*Options) {
	return func(e *Options) {
		e.useAgent = u
	}
}

func makeSigner(keyname string) (signer ssh.Signer, err error) {
	fp, err := os.Open(keyname)
	if err != nil {
		return
	}
	defer fp.Close()

	buf, _ := ioutil.ReadAll(fp)
	signer, _ = ssh.ParsePrivateKey(buf)
	return
}

func makeKeyring(key string, useAgent bool) ssh.AuthMethod {
	signers := []ssh.Signer{}

	if useAgent == true {
		aConn, _ := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
		sshAgent := agent.NewClient(aConn)
		aSigners, _ := sshAgent.Signers()
		for _, signer := range aSigners {
			signers = append(signers, signer)
		}
	}

	keys := []string{key}

	for _, keyname := range keys {
		signer, err := makeSigner(keyname)
		if err == nil {
			signers = append(signers, signer)
		}
	}
	return ssh.PublicKeys(signers...)
}

func executeCmd(opt Options, hostname string, config *ssh.ClientConfig) executeResult {

	conn, err := ssh.Dial("tcp", hostname+":"+opt.port, config)

	if err != nil {
		return executeResult{result: "",
			err: err}
	}

	session, _ := conn.NewSession()
	defer session.Close()

	var stdoutBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	err = session.Run(opt.cmd)

	return executeResult{result: hostname + ":\n" + stdoutBuf.String(),
		err: err}
}

func Run(options ...func(*Options)) bool {
	opt := Options{}
	for _, option := range options {
		option(&opt)
	}

	// in 5 seconds the message will come to timeout channel
	timeout := time.After(5 * time.Second)
	results := make(chan executeResult, len(opt.machines))

	config := &ssh.ClientConfig{
		User: opt.user,
		Auth: []ssh.AuthMethod{
			makeKeyring(opt.key, opt.useAgent),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	for _, m := range opt.machines {
		go func(hostname string) {
			// we’ll write results into the buffered channel of strings
			results <- executeCmd(opt, hostname, config)
		}(m)
	}

	for i := 0; i < len(opt.machines); i++ {
		select {
		case res := <-results:
			if res.err == nil {
				fmt.Print(res.result)
			} else {
				fmt.Println(res.err)
				return false
			}
		case <-timeout:
			fmt.Println("Timed out!")
			return false
		}
	}
	return true
}
<commit_msg>Improvement to return code and formatting (#19)<commit_after>// Copyright © 2017 Ricardo Aravena <raravena@branch.io>
//
// 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 exec

import (
	"bytes"
	"fmt"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"io/ioutil"
	"net"
	"os"
	"time"
)

type Options struct {
	machines []string
	port     string
	user     string
	cmd      string
	key      string
	useAgent bool
}

type executeResult struct {
	result string
	err    error
}

func User(u string) func(*Options) {
	return func(e *Options) {
		e.user = u
	}
}

func Port(p string) func(*Options) {
	return func(e *Options) {
		e.port = p
	}
}

func Cmd(c string) func(*Options) {
	return func(e *Options) {
		e.cmd = c
	}
}

func Machines(m []string) func(*Options) {
	return func(e *Options) {
		e.machines = m
	}
}

func Key(k string) func(*Options) {
	return func(e *Options) {
		e.key = k
	}
}

func UseAgent(u bool) func(*Options) {
	return func(e *Options) {
		e.useAgent = u
	}
}

func makeSigner(keyname string) (signer ssh.Signer, err error) {
	fp, err := os.Open(keyname)
	if err != nil {
		return
	}
	defer fp.Close()

	buf, _ := ioutil.ReadAll(fp)
	signer, _ = ssh.ParsePrivateKey(buf)
	return
}

func makeKeyring(key string, useAgent bool) ssh.AuthMethod {
	signers := []ssh.Signer{}

	if useAgent == true {
		aConn, _ := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
		sshAgent := agent.NewClient(aConn)
		aSigners, _ := sshAgent.Signers()
		for _, signer := range aSigners {
			signers = append(signers, signer)
		}
	}

	keys := []string{key}

	for _, keyname := range keys {
		signer, err := makeSigner(keyname)
		if err == nil {
			signers = append(signers, signer)
		}
	}
	return ssh.PublicKeys(signers...)
}

func executeCmd(opt Options, hostname string, config *ssh.ClientConfig) executeResult {

	conn, err := ssh.Dial("tcp", hostname+":"+opt.port, config)

	if err != nil {
		return executeResult{result: "",
			err: err}
	}

	session, _ := conn.NewSession()
	defer session.Close()

	var stdoutBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	err = session.Run(opt.cmd)

	return executeResult{result: hostname + ":\n" + stdoutBuf.String(),
		err: err}
}

func Run(options ...func(*Options)) bool {
	opt := Options{}
	for _, option := range options {
		option(&opt)
	}

	// in 5 seconds the message will come to timeout channel
	timeout := time.After(5 * time.Second)
	results := make(chan executeResult, len(opt.machines))

	config := &ssh.ClientConfig{
		User: opt.user,
		Auth: []ssh.AuthMethod{
			makeKeyring(opt.key, opt.useAgent),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	for _, m := range opt.machines {
		go func(hostname string) {
			// we’ll write results into the buffered channel of strings
			results <- executeCmd(opt, hostname, config)
		}(m)
	}

	retval := true

	for i := 0; i < len(opt.machines); i++ {
		select {
		case res := <-results:
			if res.err == nil {
				fmt.Print(res.result)
			} else {
				fmt.Println(res.err)
				retval = false
			}
		case <-timeout:
			fmt.Println(fmt.Sprintf("%v:", opt.machines[i]))
			fmt.Println("Server timed out!")
			retval = false
		}
	}
	return retval
}
<|endoftext|>


<commit_before>package rest

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sync/atomic"
	"time"

	"github.com/couchbase/sync_gateway/base"
)

var (
	// ErrSGCollectInfoAlreadyRunning is returned if sgcollect_info is already running.
	ErrSGCollectInfoAlreadyRunning = errors.New("already running")
	// ErrSGCollectInfoNotRunning is returned if sgcollect_info is not running.
	ErrSGCollectInfoNotRunning = errors.New("not running")

	sgcollectInstance = sgCollect{status: base.Uint32Ptr(sgStopped)}
)

const (
	sgStopped uint32 = iota
	sgRunning

	defaultSGUploadHost = "https://s3.amazonaws.com/cb-customers"
)

type sgCollect struct {
	cancel context.CancelFunc
	status *uint32
}

// Start will attempt to start sgcollect_info, if another is not already running.
func (sg *sgCollect) Start(zipPath string, args ...string) error {
	if atomic.LoadUint32(sg.status) == sgRunning {
		return ErrSGCollectInfoAlreadyRunning
	}

	sgPath, sgCollectPath, err := sgCollectPaths()
	if err != nil {
		return err
	}

	args = append(args, "--sync-gateway-executable", sgPath, zipPath)

	ctx, cancelFunc := context.WithCancel(context.Background())
	sg.cancel = cancelFunc
	cmd := exec.CommandContext(ctx, sgCollectPath, args...)

	// Send command stderr/stdout to pipe
	pipeReader, pipeWriter := io.Pipe()
	cmd.Stderr = pipeWriter
	cmd.Stdout = pipeWriter

	if err := cmd.Start(); err != nil {
		return err
	}

	atomic.StoreUint32(sg.status, sgRunning)
	base.Infof(base.KeyAdmin, "sgcollect_info started with args: %v", base.UD(args))

	// Stream sgcollect_info output to the logs
	go func() {
		scanner := bufio.NewScanner(pipeReader)
		for scanner.Scan() {
			base.Debugf(base.KeyAdmin, "sgcollect_info: %v", scanner.Text())
		}
		if err := scanner.Err(); err != nil {
			base.Warnf(base.KeyAdmin, "sgcollect_info: unexpected error: %v", err)
		}
	}()

	go func() {
		// Blocks until command finishes
		err := cmd.Wait()

		atomic.StoreUint32(sg.status, sgStopped)
		duration := cmd.ProcessState.UserTime()

		if err != nil {
			if err.Error() == "signal: killed" {
				base.Infof(base.KeyAdmin, "sgcollect_info cancelled after %v", duration)
				return
			}

			base.Warnf(base.KeyAdmin, "sgcollect_info failed after %v with reason: %v. Check debug logs with log key 'Admin' for more detail.", duration, err)
			return
		}

		base.Infof(base.KeyAdmin, "sgcollect_info finished successfully after %v", duration)
	}()

	return nil
}

// Stop will stop sgcollect_info, if running.
func (sg *sgCollect) Stop() error {
	if atomic.LoadUint32(sg.status) == sgStopped {
		return ErrSGCollectInfoNotRunning
	}

	sg.cancel()
	atomic.StoreUint32(sg.status, sgStopped)

	return nil
}

// IsRunning returns true if sgcollect_info is running
func (sg *sgCollect) IsRunning() bool {
	return atomic.LoadUint32(sg.status) == sgRunning
}

type sgCollectOptions struct {
	RedactLevel     string `json:"redact_level,omitempty"`
	RedactSalt      string `json:"redact_salt,omitempty"`
	OutputDirectory string `json:"output_dir,omitempty"`
	Upload          bool   `json:"upload,omitempty"`
	UploadHost      string `json:"upload_host,omitempty"`
	Customer        string `json:"customer,omitempty"`
	Ticket          string `json:"ticket,omitempty"`
}

// Validate ensures the options are OK to use in sgcollect_info.
func (c *sgCollectOptions) Validate() error {
	// Validate given output directory exists, and is a directory.
	// This does not check for write permission, however sgcollect_info
	// will fail with an error giving that reason, if this is the case.
	if c.OutputDirectory != "" {
		if fileInfo, err := os.Stat(c.OutputDirectory); err != nil {
			return err
		} else if !fileInfo.IsDir() {
			return errors.New("not a directory")
		}
	}

	if c.Upload {
		// Customer number is required if uploading.
		if c.Customer == "" {
			return errors.New("customer must be set if upload is true")
		}
		// Default uploading to support bucket if upload_host is not specified.
		if c.UploadHost == "" {
			c.UploadHost = defaultSGUploadHost
		}
	} else if c.UploadHost != "" {
		return errors.New("upload must be set to true if an upload_host is specified")
	}

	return nil
}

// Args returns a set of arguments to pass to sgcollect_info.
func (c *sgCollectOptions) Args() []string {
	var args = make([]string, 0)

	if c.Upload {
		args = append(args, "--upload-host", c.UploadHost)
	}

	if c.Customer != "" {
		args = append(args, "--customer", c.Customer)
	}

	if c.Ticket != "" {
		args = append(args, "--ticket", c.Ticket)
	}

	if c.RedactLevel != "" {
		args = append(args, "--log-redaction-level", c.RedactLevel)
	}

	if c.RedactSalt != "" {
		args = append(args, "--log-redaction-salt", c.RedactSalt)
	}

	return args
}

// sgCollectPaths attempts to return the absolute paths to Sync Gateway and to sgcollect_info binaries.
func sgCollectPaths() (sgBinary, sgCollectBinary string, err error) {
	sgBinary, err = os.Executable()
	if err != nil {
		return "", "", err
	}

	sgBinary, err = filepath.Abs(sgBinary)
	if err != nil {
		return "", "", err
	}

	hasBinDir := true
	sgCollectPath := filepath.Join("tools", "sgcollect_info")

	if runtime.GOOS == "windows" {
		sgCollectPath += ".exe"
		// Windows has no bin directory for the SG executable.
		hasBinDir = false
	}

	for {
		if hasBinDir {
			sgCollectBinary = filepath.Join(filepath.Dir(filepath.Dir(sgBinary)), sgCollectPath)
		} else {
			sgCollectBinary = filepath.Join(filepath.Dir(sgBinary), sgCollectPath)
		}

		// Check sgcollect_info exists at the path we guessed.
		base.Debugf(base.KeyAdmin, "Checking sgcollect_info binary exists at: %v", sgCollectBinary)
		_, err = os.Stat(sgCollectBinary)
		if err != nil {

			// First attempt may fail if there's no bin directory, so we'll try once more without.
			if hasBinDir {
				hasBinDir = false
				continue
			}

			return "", "", err
		}

		return sgBinary, sgCollectBinary, nil
	}
}

// sgcollectFilename returns a Windows-safe filename for sgcollect_info zip files.
func sgcollectFilename() string {

	// get timestamp
	timestamp := time.Now().UTC().Format("2006-01-02t150405")

	// use a shortened product name
	name := "sg"
	if base.ProductName == "Couchbase SG Accel" {
		name = "sga"
	}

	// get primary IP address
	ip, err := base.FindPrimaryAddr()
	if err != nil {
		ip = net.IPv4zero
	}

	// E.g: sgcollectinfo-2018-05-10t133456-sg@203.0.113.123.zip
	filename := fmt.Sprintf("sgcollectinfo-%s-%s@%s.zip", timestamp, name, ip)

	// Strip illegal Windows filename characters
	filename = base.ReplaceAll(filename, "\\/:*?\"<>|", "")

	return filename
}
<commit_msg>Clean given sgcollectoptions filepath (#3582)<commit_after>package rest

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sync/atomic"
	"time"

	"github.com/couchbase/sync_gateway/base"
	"github.com/pkg/errors"
)

var (
	// ErrSGCollectInfoAlreadyRunning is returned if sgcollect_info is already running.
	ErrSGCollectInfoAlreadyRunning = errors.New("already running")
	// ErrSGCollectInfoNotRunning is returned if sgcollect_info is not running.
	ErrSGCollectInfoNotRunning = errors.New("not running")

	sgcollectInstance = sgCollect{status: base.Uint32Ptr(sgStopped)}
)

const (
	sgStopped uint32 = iota
	sgRunning

	defaultSGUploadHost = "https://s3.amazonaws.com/cb-customers"
)

type sgCollect struct {
	cancel context.CancelFunc
	status *uint32
}

// Start will attempt to start sgcollect_info, if another is not already running.
func (sg *sgCollect) Start(zipPath string, args ...string) error {
	if atomic.LoadUint32(sg.status) == sgRunning {
		return ErrSGCollectInfoAlreadyRunning
	}

	sgPath, sgCollectPath, err := sgCollectPaths()
	if err != nil {
		return err
	}

	args = append(args, "--sync-gateway-executable", sgPath, zipPath)

	ctx, cancelFunc := context.WithCancel(context.Background())
	sg.cancel = cancelFunc
	cmd := exec.CommandContext(ctx, sgCollectPath, args...)

	// Send command stderr/stdout to pipe
	pipeReader, pipeWriter := io.Pipe()
	cmd.Stderr = pipeWriter
	cmd.Stdout = pipeWriter

	if err := cmd.Start(); err != nil {
		return err
	}

	atomic.StoreUint32(sg.status, sgRunning)
	base.Infof(base.KeyAdmin, "sgcollect_info started with args: %v", base.UD(args))

	// Stream sgcollect_info output to the logs
	go func() {
		scanner := bufio.NewScanner(pipeReader)
		for scanner.Scan() {
			base.Debugf(base.KeyAdmin, "sgcollect_info: %v", scanner.Text())
		}
		if err := scanner.Err(); err != nil {
			base.Warnf(base.KeyAdmin, "sgcollect_info: unexpected error: %v", err)
		}
	}()

	go func() {
		// Blocks until command finishes
		err := cmd.Wait()

		atomic.StoreUint32(sg.status, sgStopped)
		duration := cmd.ProcessState.UserTime()

		if err != nil {
			if err.Error() == "signal: killed" {
				base.Infof(base.KeyAdmin, "sgcollect_info cancelled after %v", duration)
				return
			}

			base.Warnf(base.KeyAdmin, "sgcollect_info failed after %v with reason: %v. Check debug logs with log key 'Admin' for more detail.", duration, err)
			return
		}

		base.Infof(base.KeyAdmin, "sgcollect_info finished successfully after %v", duration)
	}()

	return nil
}

// Stop will stop sgcollect_info, if running.
func (sg *sgCollect) Stop() error {
	if atomic.LoadUint32(sg.status) == sgStopped {
		return ErrSGCollectInfoNotRunning
	}

	sg.cancel()
	atomic.StoreUint32(sg.status, sgStopped)

	return nil
}

// IsRunning returns true if sgcollect_info is running
func (sg *sgCollect) IsRunning() bool {
	return atomic.LoadUint32(sg.status) == sgRunning
}

type sgCollectOptions struct {
	RedactLevel     string `json:"redact_level,omitempty"`
	RedactSalt      string `json:"redact_salt,omitempty"`
	OutputDirectory string `json:"output_dir,omitempty"`
	Upload          bool   `json:"upload,omitempty"`
	UploadHost      string `json:"upload_host,omitempty"`
	Customer        string `json:"customer,omitempty"`
	Ticket          string `json:"ticket,omitempty"`
}

// Validate ensures the options are OK to use in sgcollect_info.
func (c *sgCollectOptions) Validate() error {
	// Validate given output directory exists, and is a directory.
	// This does not check for write permission, however sgcollect_info
	// will fail with an error giving that reason, if this is the case.
	if c.OutputDirectory != "" {
		// Clean the given path first, for cross-platform paths.
		c.OutputDirectory = filepath.Clean(c.OutputDirectory)
		if fileInfo, err := os.Stat(c.OutputDirectory); err != nil {
			if os.IsNotExist(err) {
				return errors.Wrap(err, "no such file or directory")
			}
			return err
		} else if !fileInfo.IsDir() {
			return errors.New("not a directory")
		}
	}

	if c.Upload {
		// Customer number is required if uploading.
		if c.Customer == "" {
			return errors.New("customer must be set if upload is true")
		}
		// Default uploading to support bucket if upload_host is not specified.
		if c.UploadHost == "" {
			c.UploadHost = defaultSGUploadHost
		}
	} else if c.UploadHost != "" {
		return errors.New("upload must be set to true if an upload_host is specified")
	}

	return nil
}

// Args returns a set of arguments to pass to sgcollect_info.
func (c *sgCollectOptions) Args() []string {
	var args = make([]string, 0)

	if c.Upload {
		args = append(args, "--upload-host", c.UploadHost)
	}

	if c.Customer != "" {
		args = append(args, "--customer", c.Customer)
	}

	if c.Ticket != "" {
		args = append(args, "--ticket", c.Ticket)
	}

	if c.RedactLevel != "" {
		args = append(args, "--log-redaction-level", c.RedactLevel)
	}

	if c.RedactSalt != "" {
		args = append(args, "--log-redaction-salt", c.RedactSalt)
	}

	return args
}

// sgCollectPaths attempts to return the absolute paths to Sync Gateway and to sgcollect_info binaries.
func sgCollectPaths() (sgBinary, sgCollectBinary string, err error) {
	sgBinary, err = os.Executable()
	if err != nil {
		return "", "", err
	}

	sgBinary, err = filepath.Abs(sgBinary)
	if err != nil {
		return "", "", err
	}

	hasBinDir := true
	sgCollectPath := filepath.Join("tools", "sgcollect_info")

	if runtime.GOOS == "windows" {
		sgCollectPath += ".exe"
		// Windows has no bin directory for the SG executable.
		hasBinDir = false
	}

	for {
		if hasBinDir {
			sgCollectBinary = filepath.Join(filepath.Dir(filepath.Dir(sgBinary)), sgCollectPath)
		} else {
			sgCollectBinary = filepath.Join(filepath.Dir(sgBinary), sgCollectPath)
		}

		// Check sgcollect_info exists at the path we guessed.
		base.Debugf(base.KeyAdmin, "Checking sgcollect_info binary exists at: %v", sgCollectBinary)
		_, err = os.Stat(sgCollectBinary)
		if err != nil {

			// First attempt may fail if there's no bin directory, so we'll try once more without.
			if hasBinDir {
				hasBinDir = false
				continue
			}

			return "", "", err
		}

		return sgBinary, sgCollectBinary, nil
	}
}

// sgcollectFilename returns a Windows-safe filename for sgcollect_info zip files.
func sgcollectFilename() string {

	// get timestamp
	timestamp := time.Now().UTC().Format("2006-01-02t150405")

	// use a shortened product name
	name := "sg"
	if base.ProductName == "Couchbase SG Accel" {
		name = "sga"
	}

	// get primary IP address
	ip, err := base.FindPrimaryAddr()
	if err != nil {
		ip = net.IPv4zero
	}

	// E.g: sgcollectinfo-2018-05-10t133456-sg@203.0.113.123.zip
	filename := fmt.Sprintf("sgcollectinfo-%s-%s@%s.zip", timestamp, name, ip)

	// Strip illegal Windows filename characters
	filename = base.ReplaceAll(filename, "\\/:*?\"<>|", "")

	return filename
}
<|endoftext|>


<commit_before>package dht

import (
	"context"
	"math/rand"
	"testing"
	"time"

	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/peerstore"
	"github.com/libp2p/go-libp2p-core/protocol"
	"github.com/libp2p/go-libp2p-core/routing"

	pb "github.com/libp2p/go-libp2p-kad-dht/pb"
	record "github.com/libp2p/go-libp2p-record"
	swarmt "github.com/libp2p/go-libp2p-swarm/testing"
	bhost "github.com/libp2p/go-libp2p/p2p/host/basic"
	mocknet "github.com/libp2p/go-libp2p/p2p/net/mock"

	ggio "github.com/gogo/protobuf/io"
	u "github.com/ipfs/go-ipfs-util"
	"github.com/stretchr/testify/require"
)

// Test that one hung request to a peer doesn't prevent another request
// using that same peer from obeying its context.
func TestHungRequest(t *testing.T) {
	ctx := context.Background()
	mn, err := mocknet.FullMeshConnected(ctx, 2)
	if err != nil {
		t.Fatal(err)
	}
	hosts := mn.Hosts()

	os := []Option{testPrefix, DisableAutoRefresh()}
	d, err := New(ctx, hosts[0], os...)
	if err != nil {
		t.Fatal(err)
	}
	for _, proto := range d.serverProtocols {
		// Hang on every request.
		hosts[1].SetStreamHandler(proto, func(s network.Stream) {
			defer s.Reset() //nolint
			<-ctx.Done()
		})
	}

	require.NoError(t, hosts[0].Peerstore().AddProtocols(hosts[1].ID(), protocol.ConvertToStrings(d.serverProtocols)...))
	d.peerFound(ctx, hosts[1].ID(), true)

	ctx1, cancel1 := context.WithTimeout(ctx, 1*time.Second)
	defer cancel1()

	done := make(chan error, 1)
	go func() {
		_, err := d.GetClosestPeers(ctx1, testCaseCids[0].KeyString())
		done <- err
	}()

	time.Sleep(100 * time.Millisecond)
	ctx2, cancel2 := context.WithTimeout(ctx, 100*time.Millisecond)
	defer cancel2()
	err = d.Provide(ctx2, testCaseCids[0], true)
	if err != context.DeadlineExceeded {
		t.Errorf("expected to fail with deadline exceeded, got: %s", ctx2.Err())
	}
	select {
	case <-done:
		t.Errorf("GetClosestPeers should not have returned yet")
	default:
		err = <-done
		if err != context.DeadlineExceeded {
			t.Errorf("expected the deadline to be exceeded, got %s", err)
		}
	}

}

func TestGetFailures(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := context.Background()

	host1 := bhost.New(swarmt.GenSwarm(t, ctx, swarmt.OptDisableReuseport))
	host2 := bhost.New(swarmt.GenSwarm(t, ctx, swarmt.OptDisableReuseport))

	d, err := New(ctx, host1, testPrefix, DisableAutoRefresh(), Mode(ModeServer))
	if err != nil {
		t.Fatal(err)
	}

	// Reply with failures to every message
	for _, proto := range d.serverProtocols {
		host2.SetStreamHandler(proto, func(s network.Stream) {
			time.Sleep(400 * time.Millisecond)
			s.Close()
		})
	}

	host1.Peerstore().AddAddrs(host2.ID(), host2.Addrs(), peerstore.ConnectedAddrTTL)
	_, err = host1.Network().DialPeer(ctx, host2.ID())
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(1 * time.Second)

	// This one should time out
	ctx1, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
	defer cancel()
	if _, err := d.GetValue(ctx1, "test"); err != nil {
		if merr, ok := err.(u.MultiErr); ok && len(merr) > 0 {
			err = merr[0]
		}

		if err != context.DeadlineExceeded {
			t.Fatal("Got different error than we expected", err)
		}
	} else {
		t.Fatal("Did not get expected error!")
	}

	t.Log("Timeout test passed.")

	for _, proto := range d.serverProtocols {
		// Reply with failures to every message
		host2.SetStreamHandler(proto, func(s network.Stream) {
			defer s.Close()

			pbr := ggio.NewDelimitedReader(s, network.MessageSizeMax)
			pbw := ggio.NewDelimitedWriter(s)

			pmes := new(pb.Message)
			if err := pbr.ReadMsg(pmes); err != nil {
				// user gave up
				return
			}

			resp := &pb.Message{
				Type: pmes.Type,
			}
			_ = pbw.WriteMsg(resp)
		})
	}

	// This one should fail with NotFound.
	// long context timeout to ensure we dont end too early.
	// the dht should be exhausting its query and returning not found.
	// (was 3 seconds before which should be _plenty_ of time, but maybe
	// travis machines really have a hard time...)
	ctx2, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	_, err = d.GetValue(ctx2, "test")
	if err != nil {
		if merr, ok := err.(u.MultiErr); ok && len(merr) > 0 {
			err = merr[0]
		}
		if err != routing.ErrNotFound {
			t.Fatalf("Expected ErrNotFound, got: %s", err)
		}
	} else {
		t.Fatal("expected error, got none.")
	}

	t.Log("ErrNotFound check passed!")

	// Now we test this DHT's handleGetValue failure
	{
		typ := pb.Message_GET_VALUE
		str := "hello"

		rec := record.MakePutRecord(str, []byte("blah"))
		req := pb.Message{
			Type:   typ,
			Key:    []byte(str),
			Record: rec,
		}

		s, err := host2.NewStream(context.Background(), host1.ID(), d.protocols...)
		if err != nil {
			t.Fatal(err)
		}
		defer s.Close()

		pbr := ggio.NewDelimitedReader(s, network.MessageSizeMax)
		pbw := ggio.NewDelimitedWriter(s)

		if err := pbw.WriteMsg(&req); err != nil {
			t.Fatal(err)
		}

		pmes := new(pb.Message)
		if err := pbr.ReadMsg(pmes); err != nil {
			t.Fatal(err)
		}
		if pmes.GetRecord() != nil {
			t.Fatal("shouldnt have value")
		}
		if pmes.GetProviderPeers() != nil {
			t.Fatal("shouldnt have provider peers")
		}
	}
}

func TestNotFound(t *testing.T) {
	// t.Skip("skipping test to debug another")
	if testing.Short() {
		t.SkipNow()
	}

	ctx := context.Background()
	mn, err := mocknet.FullMeshConnected(ctx, 16)
	if err != nil {
		t.Fatal(err)
	}
	hosts := mn.Hosts()

	os := []Option{testPrefix, DisableAutoRefresh()}
	d, err := New(ctx, hosts[0], os...)
	if err != nil {
		t.Fatal(err)
	}

	for _, p := range hosts {
		d.peerFound(ctx, p.ID(), true)
	}

	// Reply with random peers to every message
	for _, host := range hosts {
		host := host // shadow loop var
		for _, proto := range d.serverProtocols {
			host.SetStreamHandler(proto, func(s network.Stream) {
				defer s.Close()

				pbr := ggio.NewDelimitedReader(s, network.MessageSizeMax)
				pbw := ggio.NewDelimitedWriter(s)

				pmes := new(pb.Message)
				if err := pbr.ReadMsg(pmes); err != nil {
					panic(err)
				}

				switch pmes.GetType() {
				case pb.Message_GET_VALUE:
					resp := &pb.Message{Type: pmes.Type}

					ps := []peer.AddrInfo{}
					for i := 0; i < 7; i++ {
						p := hosts[rand.Intn(len(hosts))].ID()
						pi := host.Peerstore().PeerInfo(p)
						ps = append(ps, pi)
					}

					resp.CloserPeers = pb.PeerInfosToPBPeers(d.host.Network(), ps)
					if err := pbw.WriteMsg(resp); err != nil {
						panic(err)
					}
				default:
					panic("Shouldnt recieve this.")
				}
			})
		}
	}

	// long timeout to ensure timing is not at play.
	ctx, cancel := context.WithTimeout(ctx, time.Second*20)
	defer cancel()
	v, err := d.GetValue(ctx, "hello")
	logger.Debugf("get value got %v", v)
	if err != nil {
		if merr, ok := err.(u.MultiErr); ok && len(merr) > 0 {
			err = merr[0]
		}
		switch err {
		case routing.ErrNotFound:
			//Success!
			return
		case u.ErrTimeout:
			t.Fatal("Should not have gotten timeout!")
		default:
			t.Fatalf("Got unexpected error: %s", err)
		}
	}
	t.Fatal("Expected to recieve an error.")
}

// If less than K nodes are in the entire network, it should fail when we make
// a GET rpc and nobody has the value
func TestLessThanKResponses(t *testing.T) {
	// t.Skip("skipping test to debug another")
	// t.Skip("skipping test because it makes a lot of output")

	ctx := context.Background()
	mn, err := mocknet.FullMeshConnected(ctx, 6)
	if err != nil {
		t.Fatal(err)
	}
	hosts := mn.Hosts()

	os := []Option{testPrefix, DisableAutoRefresh()}
	d, err := New(ctx, hosts[0], os...)
	if err != nil {
		t.Fatal(err)
	}

	for i := 1; i < 5; i++ {
		d.peerFound(ctx, hosts[i].ID(), true)
	}

	// Reply with random peers to every message
	for _, host := range hosts {
		host := host // shadow loop var
		for _, proto := range d.serverProtocols {
			host.SetStreamHandler(proto, func(s network.Stream) {
				defer s.Close()

				pbr := ggio.NewDelimitedReader(s, network.MessageSizeMax)
				pbw := ggio.NewDelimitedWriter(s)

				pmes := new(pb.Message)
				if err := pbr.ReadMsg(pmes); err != nil {
					panic(err)
				}

				switch pmes.GetType() {
				case pb.Message_GET_VALUE:
					pi := host.Peerstore().PeerInfo(hosts[1].ID())
					resp := &pb.Message{
						Type:        pmes.Type,
						CloserPeers: pb.PeerInfosToPBPeers(d.host.Network(), []peer.AddrInfo{pi}),
					}

					if err := pbw.WriteMsg(resp); err != nil {
						panic(err)
					}
				default:
					panic("Shouldnt recieve this.")
				}

			})
		}
	}

	ctx, cancel := context.WithTimeout(ctx, time.Second*30)
	defer cancel()
	if _, err := d.GetValue(ctx, "hello"); err != nil {
		switch err {
		case routing.ErrNotFound:
			//Success!
			return
		case u.ErrTimeout:
			t.Fatal("Should not have gotten timeout!")
		default:
			t.Fatalf("Got unexpected error: %s", err)
		}
	}
	t.Fatal("Expected to recieve an error.")
}

// Test multiple queries against a node that closes its stream after every query.
func TestMultipleQueries(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := context.Background()
	mn, err := mocknet.FullMeshConnected(ctx, 2)
	if err != nil {
		t.Fatal(err)
	}
	hosts := mn.Hosts()
	os := []Option{testPrefix, DisableAutoRefresh()}
	d, err := New(ctx, hosts[0], os...)
	if err != nil {
		t.Fatal(err)
	}

	d.peerFound(ctx, hosts[1].ID(), true)

	for _, proto := range d.serverProtocols {
		// It would be nice to be able to just get a value and succeed but then
		// we'd need to deal with selectors and validators...
		hosts[1].SetStreamHandler(proto, func(s network.Stream) {
			defer s.Close()

			pbr := ggio.NewDelimitedReader(s, network.MessageSizeMax)
			pbw := ggio.NewDelimitedWriter(s)

			pmes := new(pb.Message)
			if err := pbr.ReadMsg(pmes); err != nil {
				panic(err)
			}

			switch pmes.GetType() {
			case pb.Message_GET_VALUE:
				pi := hosts[1].Peerstore().PeerInfo(hosts[0].ID())
				resp := &pb.Message{
					Type:        pmes.Type,
					CloserPeers: pb.PeerInfosToPBPeers(d.host.Network(), []peer.AddrInfo{pi}),
				}

				if err := pbw.WriteMsg(resp); err != nil {
					panic(err)
				}
			default:
				panic("Shouldnt recieve this.")
			}
		})
	}

	// long timeout to ensure timing is not at play.
	ctx, cancel := context.WithTimeout(ctx, time.Second*20)
	defer cancel()
	for i := 0; i < 10; i++ {
		if _, err := d.GetValue(ctx, "hello"); err != nil {
			switch err {
			case routing.ErrNotFound:
				//Success!
				continue
			case u.ErrTimeout:
				t.Fatal("Should not have gotten timeout!")
			default:
				t.Fatalf("Got unexpected error: %s", err)
			}
		}
		t.Fatal("Expected to recieve an error.")
	}
}
<commit_msg>test: use server-mode everywhere<commit_after>package dht

import (
	"context"
	"math/rand"
	"testing"
	"time"

	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/peerstore"
	"github.com/libp2p/go-libp2p-core/protocol"
	"github.com/libp2p/go-libp2p-core/routing"

	pb "github.com/libp2p/go-libp2p-kad-dht/pb"
	record "github.com/libp2p/go-libp2p-record"
	swarmt "github.com/libp2p/go-libp2p-swarm/testing"
	bhost "github.com/libp2p/go-libp2p/p2p/host/basic"
	mocknet "github.com/libp2p/go-libp2p/p2p/net/mock"

	ggio "github.com/gogo/protobuf/io"
	u "github.com/ipfs/go-ipfs-util"
	"github.com/stretchr/testify/require"
)

// Test that one hung request to a peer doesn't prevent another request
// using that same peer from obeying its context.
func TestHungRequest(t *testing.T) {
	ctx := context.Background()
	mn, err := mocknet.FullMeshConnected(ctx, 2)
	if err != nil {
		t.Fatal(err)
	}
	hosts := mn.Hosts()

	os := []Option{testPrefix, DisableAutoRefresh(), Mode(ModeServer)}
	d, err := New(ctx, hosts[0], os...)
	if err != nil {
		t.Fatal(err)
	}
	for _, proto := range d.serverProtocols {
		// Hang on every request.
		hosts[1].SetStreamHandler(proto, func(s network.Stream) {
			defer s.Reset() //nolint
			<-ctx.Done()
		})
	}

	require.NoError(t, hosts[0].Peerstore().AddProtocols(hosts[1].ID(), protocol.ConvertToStrings(d.serverProtocols)...))
	d.peerFound(ctx, hosts[1].ID(), true)

	ctx1, cancel1 := context.WithTimeout(ctx, 1*time.Second)
	defer cancel1()

	done := make(chan error, 1)
	go func() {
		_, err := d.GetClosestPeers(ctx1, testCaseCids[0].KeyString())
		done <- err
	}()

	time.Sleep(100 * time.Millisecond)
	ctx2, cancel2 := context.WithTimeout(ctx, 100*time.Millisecond)
	defer cancel2()
	err = d.Provide(ctx2, testCaseCids[0], true)
	if err != context.DeadlineExceeded {
		t.Errorf("expected to fail with deadline exceeded, got: %s", ctx2.Err())
	}
	select {
	case <-done:
		t.Errorf("GetClosestPeers should not have returned yet")
	default:
		err = <-done
		if err != context.DeadlineExceeded {
			t.Errorf("expected the deadline to be exceeded, got %s", err)
		}
	}

}

func TestGetFailures(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := context.Background()

	host1 := bhost.New(swarmt.GenSwarm(t, ctx, swarmt.OptDisableReuseport))
	host2 := bhost.New(swarmt.GenSwarm(t, ctx, swarmt.OptDisableReuseport))

	d, err := New(ctx, host1, testPrefix, DisableAutoRefresh(), Mode(ModeServer))
	if err != nil {
		t.Fatal(err)
	}

	// Reply with failures to every message
	for _, proto := range d.serverProtocols {
		host2.SetStreamHandler(proto, func(s network.Stream) {
			time.Sleep(400 * time.Millisecond)
			s.Close()
		})
	}

	host1.Peerstore().AddAddrs(host2.ID(), host2.Addrs(), peerstore.ConnectedAddrTTL)
	_, err = host1.Network().DialPeer(ctx, host2.ID())
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(1 * time.Second)

	// This one should time out
	ctx1, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
	defer cancel()
	if _, err := d.GetValue(ctx1, "test"); err != nil {
		if merr, ok := err.(u.MultiErr); ok && len(merr) > 0 {
			err = merr[0]
		}

		if err != context.DeadlineExceeded {
			t.Fatal("Got different error than we expected", err)
		}
	} else {
		t.Fatal("Did not get expected error!")
	}

	t.Log("Timeout test passed.")

	for _, proto := range d.serverProtocols {
		// Reply with failures to every message
		host2.SetStreamHandler(proto, func(s network.Stream) {
			defer s.Close()

			pbr := ggio.NewDelimitedReader(s, network.MessageSizeMax)
			pbw := ggio.NewDelimitedWriter(s)

			pmes := new(pb.Message)
			if err := pbr.ReadMsg(pmes); err != nil {
				// user gave up
				return
			}

			resp := &pb.Message{
				Type: pmes.Type,
			}
			_ = pbw.WriteMsg(resp)
		})
	}

	// This one should fail with NotFound.
	// long context timeout to ensure we dont end too early.
	// the dht should be exhausting its query and returning not found.
	// (was 3 seconds before which should be _plenty_ of time, but maybe
	// travis machines really have a hard time...)
	ctx2, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	_, err = d.GetValue(ctx2, "test")
	if err != nil {
		if merr, ok := err.(u.MultiErr); ok && len(merr) > 0 {
			err = merr[0]
		}
		if err != routing.ErrNotFound {
			t.Fatalf("Expected ErrNotFound, got: %s", err)
		}
	} else {
		t.Fatal("expected error, got none.")
	}

	t.Log("ErrNotFound check passed!")

	// Now we test this DHT's handleGetValue failure
	{
		typ := pb.Message_GET_VALUE
		str := "hello"

		rec := record.MakePutRecord(str, []byte("blah"))
		req := pb.Message{
			Type:   typ,
			Key:    []byte(str),
			Record: rec,
		}

		s, err := host2.NewStream(context.Background(), host1.ID(), d.protocols...)
		if err != nil {
			t.Fatal(err)
		}
		defer s.Close()

		pbr := ggio.NewDelimitedReader(s, network.MessageSizeMax)
		pbw := ggio.NewDelimitedWriter(s)

		if err := pbw.WriteMsg(&req); err != nil {
			t.Fatal(err)
		}

		pmes := new(pb.Message)
		if err := pbr.ReadMsg(pmes); err != nil {
			t.Fatal(err)
		}
		if pmes.GetRecord() != nil {
			t.Fatal("shouldnt have value")
		}
		if pmes.GetProviderPeers() != nil {
			t.Fatal("shouldnt have provider peers")
		}
	}
}

func TestNotFound(t *testing.T) {
	// t.Skip("skipping test to debug another")
	if testing.Short() {
		t.SkipNow()
	}

	ctx := context.Background()
	mn, err := mocknet.FullMeshConnected(ctx, 16)
	if err != nil {
		t.Fatal(err)
	}
	hosts := mn.Hosts()

	os := []Option{testPrefix, DisableAutoRefresh(), Mode(ModeServer)}
	d, err := New(ctx, hosts[0], os...)
	if err != nil {
		t.Fatal(err)
	}

	for _, p := range hosts {
		d.peerFound(ctx, p.ID(), true)
	}

	// Reply with random peers to every message
	for _, host := range hosts {
		host := host // shadow loop var
		for _, proto := range d.serverProtocols {
			host.SetStreamHandler(proto, func(s network.Stream) {
				defer s.Close()

				pbr := ggio.NewDelimitedReader(s, network.MessageSizeMax)
				pbw := ggio.NewDelimitedWriter(s)

				pmes := new(pb.Message)
				if err := pbr.ReadMsg(pmes); err != nil {
					panic(err)
				}

				switch pmes.GetType() {
				case pb.Message_GET_VALUE:
					resp := &pb.Message{Type: pmes.Type}

					ps := []peer.AddrInfo{}
					for i := 0; i < 7; i++ {
						p := hosts[rand.Intn(len(hosts))].ID()
						pi := host.Peerstore().PeerInfo(p)
						ps = append(ps, pi)
					}

					resp.CloserPeers = pb.PeerInfosToPBPeers(d.host.Network(), ps)
					if err := pbw.WriteMsg(resp); err != nil {
						panic(err)
					}
				default:
					panic("Shouldnt recieve this.")
				}
			})
		}
	}

	// long timeout to ensure timing is not at play.
	ctx, cancel := context.WithTimeout(ctx, time.Second*20)
	defer cancel()
	v, err := d.GetValue(ctx, "hello")
	logger.Debugf("get value got %v", v)
	if err != nil {
		if merr, ok := err.(u.MultiErr); ok && len(merr) > 0 {
			err = merr[0]
		}
		switch err {
		case routing.ErrNotFound:
			//Success!
			return
		case u.ErrTimeout:
			t.Fatal("Should not have gotten timeout!")
		default:
			t.Fatalf("Got unexpected error: %s", err)
		}
	}
	t.Fatal("Expected to recieve an error.")
}

// If less than K nodes are in the entire network, it should fail when we make
// a GET rpc and nobody has the value
func TestLessThanKResponses(t *testing.T) {
	// t.Skip("skipping test to debug another")
	// t.Skip("skipping test because it makes a lot of output")

	ctx := context.Background()
	mn, err := mocknet.FullMeshConnected(ctx, 6)
	if err != nil {
		t.Fatal(err)
	}
	hosts := mn.Hosts()

	os := []Option{testPrefix, DisableAutoRefresh(), Mode(ModeServer)}
	d, err := New(ctx, hosts[0], os...)
	if err != nil {
		t.Fatal(err)
	}

	for i := 1; i < 5; i++ {
		d.peerFound(ctx, hosts[i].ID(), true)
	}

	// Reply with random peers to every message
	for _, host := range hosts {
		host := host // shadow loop var
		for _, proto := range d.serverProtocols {
			host.SetStreamHandler(proto, func(s network.Stream) {
				defer s.Close()

				pbr := ggio.NewDelimitedReader(s, network.MessageSizeMax)
				pbw := ggio.NewDelimitedWriter(s)

				pmes := new(pb.Message)
				if err := pbr.ReadMsg(pmes); err != nil {
					panic(err)
				}

				switch pmes.GetType() {
				case pb.Message_GET_VALUE:
					pi := host.Peerstore().PeerInfo(hosts[1].ID())
					resp := &pb.Message{
						Type:        pmes.Type,
						CloserPeers: pb.PeerInfosToPBPeers(d.host.Network(), []peer.AddrInfo{pi}),
					}

					if err := pbw.WriteMsg(resp); err != nil {
						panic(err)
					}
				default:
					panic("Shouldnt recieve this.")
				}

			})
		}
	}

	ctx, cancel := context.WithTimeout(ctx, time.Second*30)
	defer cancel()
	if _, err := d.GetValue(ctx, "hello"); err != nil {
		switch err {
		case routing.ErrNotFound:
			//Success!
			return
		case u.ErrTimeout:
			t.Fatal("Should not have gotten timeout!")
		default:
			t.Fatalf("Got unexpected error: %s", err)
		}
	}
	t.Fatal("Expected to recieve an error.")
}

// Test multiple queries against a node that closes its stream after every query.
func TestMultipleQueries(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := context.Background()
	mn, err := mocknet.FullMeshConnected(ctx, 2)
	if err != nil {
		t.Fatal(err)
	}
	hosts := mn.Hosts()
	os := []Option{testPrefix, DisableAutoRefresh(), Mode(ModeServer)}
	d, err := New(ctx, hosts[0], os...)
	if err != nil {
		t.Fatal(err)
	}

	d.peerFound(ctx, hosts[1].ID(), true)

	for _, proto := range d.serverProtocols {
		// It would be nice to be able to just get a value and succeed but then
		// we'd need to deal with selectors and validators...
		hosts[1].SetStreamHandler(proto, func(s network.Stream) {
			defer s.Close()

			pbr := ggio.NewDelimitedReader(s, network.MessageSizeMax)
			pbw := ggio.NewDelimitedWriter(s)

			pmes := new(pb.Message)
			if err := pbr.ReadMsg(pmes); err != nil {
				panic(err)
			}

			switch pmes.GetType() {
			case pb.Message_GET_VALUE:
				pi := hosts[1].Peerstore().PeerInfo(hosts[0].ID())
				resp := &pb.Message{
					Type:        pmes.Type,
					CloserPeers: pb.PeerInfosToPBPeers(d.host.Network(), []peer.AddrInfo{pi}),
				}

				if err := pbw.WriteMsg(resp); err != nil {
					panic(err)
				}
			default:
				panic("Shouldnt recieve this.")
			}
		})
	}

	// long timeout to ensure timing is not at play.
	ctx, cancel := context.WithTimeout(ctx, time.Second*20)
	defer cancel()
	for i := 0; i < 10; i++ {
		if _, err := d.GetValue(ctx, "hello"); err != nil {
			switch err {
			case routing.ErrNotFound:
				//Success!
				continue
			case u.ErrTimeout:
				t.Fatal("Should not have gotten timeout!")
			default:
				t.Fatalf("Got unexpected error: %s", err)
			}
		}
		t.Fatal("Expected to recieve an error.")
	}
}
<|endoftext|>


<commit_before>// +build darwin dragonfly freebsd linux netbsd openbsd plan9

package daemon

import (
	"syscall"
)

func lockFile(fd uintptr) error {
	err := syscall.Flock(int(fd), syscall.LOCK_EX|syscall.LOCK_NB)
	if err == syscall.EWOULDBLOCK {
		err = ErrWouldBlock
	}
	return err
}

func unlockFile(fd uintptr) error {
	err := syscall.Flock(int(fd), syscall.LOCK_UN)
	if err == syscall.EWOULDBLOCK {
		err = ErrWouldBlock
	}
	return err
}
<commit_msg>Fixed: lock_file_unix requires to ignore solaris<commit_after>// +build darwin dragonfly freebsd linux netbsd openbsd plan9 !solaris

package daemon

import (
	"syscall"
)

func lockFile(fd uintptr) error {
	err := syscall.Flock(int(fd), syscall.LOCK_EX|syscall.LOCK_NB)
	if err == syscall.EWOULDBLOCK {
		err = ErrWouldBlock
	}
	return err
}

func unlockFile(fd uintptr) error {
	err := syscall.Flock(int(fd), syscall.LOCK_UN)
	if err == syscall.EWOULDBLOCK {
		err = ErrWouldBlock
	}
	return err
}
<|endoftext|>