<commit_before>package shells

import (
	"path"
	"path/filepath"
	"strconv"
	"strings"

	"errors"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
)

type AbstractShell struct {
}

func (b *AbstractShell) GetFeatures(features *common.FeaturesInfo) {
	features.Artifacts = true
	features.Cache = true
}

func (b *AbstractShell) GetSupportedOptions() []string {
	return []string{"artifacts", "cache", "dependencies", "after_script"}
}

func (b *AbstractShell) writeCdBuildDir(w ShellWriter, info common.ShellScriptInfo) {
	w.Cd(info.Build.FullProjectDir())
}

func (b *AbstractShell) writeExports(w ShellWriter, info common.ShellScriptInfo) {
	for _, variable := range info.Build.GetAllVariables() {
		w.Variable(variable)
	}
}

func (b *AbstractShell) writeTLSCAInfo(w ShellWriter, build *common.Build, key string) {
	if build.TLSCAChain != "" {
		w.Variable(common.BuildVariable{
			Key:      key,
			Value:    build.TLSCAChain,
			Public:   true,
			Internal: true,
			File:     true,
		})
	}
}

func (b *AbstractShell) writeCloneCmd(w ShellWriter, build *common.Build, projectDir string) {
	w.RmDir(projectDir)
	if depth := build.GetGitDepth(); depth != "" {
		w.Notice("Cloning repository for %s with git depth set to %s...", build.RefName, depth)
		w.Command("git", "clone", build.RepoURL, projectDir, "--depth", depth, "--branch", build.RefName)
	} else {
		w.Notice("Cloning repository...")
		w.Command("git", "clone", build.RepoURL, projectDir)
	}
	w.Cd(projectDir)
}

func (b *AbstractShell) writeFetchCmd(w ShellWriter, build *common.Build, projectDir string, gitDir string) {
	depth := build.GetGitDepth()

	w.IfDirectory(gitDir)
	if depth != "" {
		w.Notice("Fetching changes for %s with git depth set to %s...", build.RefName, depth)
	} else {
		w.Notice("Fetching changes...")
	}
	w.Cd(projectDir)
	w.Command("git", "clean", "-ffdx")
	w.Command("git", "reset", "--hard")
	w.Command("git", "remote", "set-url", "origin", build.RepoURL)
	if depth != "" {
		var refspec string
		if build.Tag {
			refspec = "+refs/tags/" + build.RefName + ":refs/tags/" + build.RefName
		} else {
			refspec = "+refs/heads/" + build.RefName + ":refs/remotes/origin/" + build.RefName
		}
		w.Command("git", "fetch", "--depth", depth, "origin", "--prune", refspec)
	} else {
		w.Command("git", "fetch", "origin", "--prune", "+refs/heads/*:refs/remotes/origin/*", "+refs/tags/*:refs/tags/*")
	}
	w.Else()
	b.writeCloneCmd(w, build, projectDir)
	w.EndIf()
}

func (b *AbstractShell) writeCheckoutCmd(w ShellWriter, build *common.Build) {
	w.Notice("Checking out %s as %s...", build.Sha[0:8], build.RefName)
	// We remove a git index file, this is required if `git checkout` is terminated
	w.RmFile(".git/index.lock")
	w.Command("git", "checkout", "-q", build.Sha)
}

func (b *AbstractShell) cacheFile(build *common.Build, userKey string) (key, file string) {
	if build.CacheDir == "" {
		return
	}

	// Deduce cache key
	key = path.Join(build.Name, build.RefName)
	if userKey != "" {
		key = build.GetAllVariables().ExpandValue(userKey)
	}

	// Ignore cache without the key
	if key == "" {
		return
	}

	file = path.Join(build.CacheDir, key, "cache.zip")
	file, err := filepath.Rel(build.BuildDir, file)
	if err != nil {
		return "", ""
	}
	return
}

func (o *archivingOptions) CommandArguments() (args []string) {
	for _, path := range o.Paths {
		args = append(args, "--path", path)
	}

	if o.Untracked {
		args = append(args, "--untracked")
	}
	return
}

func (b *AbstractShell) guardRunnerCommand(w ShellWriter, runnerCommand string, action string, f func()) {
	if runnerCommand == "" {
		w.Warning("%s is not supported by this executor.", action)
		return
	}

	w.IfCmd(runnerCommand, "--version")
	f()
	w.Else()
	w.Warning("Missing %s. %s is disabled.", runnerCommand, action)
	w.EndIf()
}

func (b *AbstractShell) cacheExtractor(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}

	// Skip restoring cache if no cache is defined
	if archiverArgs := options.CommandArguments(); len(archiverArgs) == 0 {
		return
	}

	// Skip archiving if no cache is defined
	cacheKey, cacheFile := b.cacheFile(info.Build, options.Key)
	if cacheKey == "" {
		return
	}

	args := []string{
		"cache-extractor",
		"--file", cacheFile,
	}

	// Generate cache download address
	if url := getCacheDownloadURL(info.Build, cacheKey); url != nil {
		args = append(args, "--url", url.String())
	}

	// Execute archive command
	b.guardRunnerCommand(w, info.RunnerCommand, "Extracting cache", func() {
		w.Notice("Checking cache for %s...", cacheKey)
		w.Command(info.RunnerCommand, args...)
	})
}

func (b *AbstractShell) downloadArtifacts(w ShellWriter, build *common.BuildInfo, info common.ShellScriptInfo) {
	args := []string{
		"artifacts-downloader",
		"--url",
		info.Build.Runner.URL,
		"--token",
		build.Token,
		"--id",
		strconv.Itoa(build.ID),
	}

	w.Notice("Downloading artifacts for %s (%d)...", build.Name, build.ID)
	w.Command(info.RunnerCommand, args...)
}

func (b *AbstractShell) downloadAllArtifacts(w ShellWriter, dependencies *dependencies, info common.ShellScriptInfo) {
	b.guardRunnerCommand(w, info.RunnerCommand, "Artifacts downloading", func() {
		for _, otherBuild := range info.Build.DependsOnBuilds {
			if otherBuild.Artifacts == nil || otherBuild.Artifacts.Filename == "" {
				continue
			}
			if !dependencies.IsDependent(otherBuild.Name) {
				continue
			}
			b.downloadArtifacts(w, &otherBuild, info)
		}
	})
}

func (b *AbstractShell) writePrepareScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	b.writeExports(w, info)

	build := info.Build
	projectDir := build.FullProjectDir()
	gitDir := path.Join(build.FullProjectDir(), ".git")

	b.writeTLSCAInfo(w, info.Build, "GIT_SSL_CAINFO")
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	w.Command("git", "config", "--global", "fetch.recurseSubmodules", "false")
	switch info.Build.GetGitStrategy() {
	case common.GitFetch:
		b.writeFetchCmd(w, build, projectDir, gitDir)

	case common.GitClone:
		b.writeCloneCmd(w, build, projectDir)

	default:
		return errors.New("unknown GIT_STRATEGY")
	}

	b.writeCheckoutCmd(w, build)

	// Parse options
	var options shellOptions
	err = info.Build.Options.Decode(&options)
	if err != nil {
		return
	}

	// Try to restore from main cache, if not found cache for master
	b.cacheExtractor(w, options.Cache, info)

	// Process all artifacts
	b.downloadAllArtifacts(w, options.Dependencies, info)
	return nil
}

func (b *AbstractShell) writeBuildScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)

	commands := info.Build.Commands
	commands = strings.TrimSpace(commands)
	for _, command := range strings.Split(commands, "\n") {
		command = strings.TrimSpace(command)
		if command != "" {
			w.Notice("$ %s", command)
		} else {
			w.EmptyLine()
		}
		w.Line(command)
		w.CheckForErrors()
	}

	return nil
}

func (b *AbstractShell) cacheArchiver(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}

	// Skip archiving if no cache is defined
	cacheKey, cacheFile := b.cacheFile(info.Build, options.Key)
	if cacheKey == "" {
		return
	}

	args := []string{
		"cache-archiver",
		"--file", cacheFile,
	}

	// Create list of files to archive
	archiverArgs := options.CommandArguments()
	if len(archiverArgs) == 0 {
		// Skip creating archive
		return
	}
	args = append(args, archiverArgs...)

	// Generate cache upload address
	if url := getCacheUploadURL(info.Build, cacheKey); url != nil {
		args = append(args, "--url", url.String())
	}

	b.guardRunnerCommand(w, info.RunnerCommand, "Creating cache", func() {
		// Execute archive command
		w.Notice("Creating cache %s...", cacheKey)
		w.Command(info.RunnerCommand, args...)
	})
}

func (b *AbstractShell) uploadArtifacts(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}
	if info.Build.Runner.URL == "" {
		return
	}

	args := []string{
		"artifacts-uploader",
		"--url",
		info.Build.Runner.URL,
		"--token",
		info.Build.Token,
		"--id",
		strconv.Itoa(info.Build.ID),
	}

	// Create list of files to archive
	archiverArgs := options.CommandArguments()
	if len(archiverArgs) == 0 {
		// Skip creating archive
		return
	}
	args = append(args, archiverArgs...)

	// Get artifacts:name
	if name, ok := info.Build.Options.GetString("artifacts", "name"); ok && name != "" {
		args = append(args, "--name", name)
	}

	// Get artifacts:expire_in
	if expireIn, ok := info.Build.Options.GetString("artifacts", "expire_in"); ok && expireIn != "" {
		args = append(args, "--expire-in", expireIn)
	}

	b.guardRunnerCommand(w, info.RunnerCommand, "Uploading artifacts", func() {
		w.Notice("Uploading artifacts...")
		w.Command(info.RunnerCommand, args...)
	})
}

func (b *AbstractShell) writeAfterScript(w ShellWriter, info common.ShellScriptInfo) error {
	shellOptions := struct {
		AfterScript []string `json:"after_script"`
	}{}
	err := info.Build.Options.Decode(&shellOptions)
	if err != nil {
		return err
	}

	if len(shellOptions.AfterScript) == 0 {
		return nil
	}

	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)

	w.Notice("Running after script...")

	for _, command := range shellOptions.AfterScript {
		command = strings.TrimSpace(command)
		if command != "" {
			w.Notice("$ %s", command)
		} else {
			w.EmptyLine()
		}
		w.Line(command)
		w.CheckForErrors()
	}

	return nil
}

func (b *AbstractShell) writeArchiveCacheScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	// Parse options
	var options shellOptions
	err = info.Build.Options.Decode(&options)
	if err != nil {
		return
	}

	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	// Find cached files and archive them
	b.cacheArchiver(w, options.Cache, info)
	return
}

func (b *AbstractShell) writeUploadArtifactsScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	// Parse options
	var options shellOptions
	err = info.Build.Options.Decode(&options)
	if err != nil {
		return
	}

	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	// Upload artifacts
	b.uploadArtifacts(w, options.Artifacts, info)
	return
}

func (b *AbstractShell) writeScript(w ShellWriter, scriptType common.ShellScriptType, info common.ShellScriptInfo) (err error) {
	switch scriptType {
	case common.ShellPrepareScript:
		return b.writePrepareScript(w, info)

	case common.ShellBuildScript:
		return b.writeBuildScript(w, info)

	case common.ShellAfterScript:
		return b.writeAfterScript(w, info)

	case common.ShellArchiveCache:
		return b.writeArchiveCacheScript(w, info)

	case common.ShellUploadArtifacts:
		return b.writeUploadArtifactsScript(w, info)

	default:
		return errors.New("Not supported script type: " + string(scriptType))
	}
}
<commit_msg>Fix broken script generation for artifacts<commit_after>package shells

import (
	"path"
	"path/filepath"
	"strconv"
	"strings"

	"errors"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
)

type AbstractShell struct {
}

func (b *AbstractShell) GetFeatures(features *common.FeaturesInfo) {
	features.Artifacts = true
	features.Cache = true
}

func (b *AbstractShell) GetSupportedOptions() []string {
	return []string{"artifacts", "cache", "dependencies", "after_script"}
}

func (b *AbstractShell) writeCdBuildDir(w ShellWriter, info common.ShellScriptInfo) {
	w.Cd(info.Build.FullProjectDir())
}

func (b *AbstractShell) writeExports(w ShellWriter, info common.ShellScriptInfo) {
	for _, variable := range info.Build.GetAllVariables() {
		w.Variable(variable)
	}
}

func (b *AbstractShell) writeTLSCAInfo(w ShellWriter, build *common.Build, key string) {
	if build.TLSCAChain != "" {
		w.Variable(common.BuildVariable{
			Key:      key,
			Value:    build.TLSCAChain,
			Public:   true,
			Internal: true,
			File:     true,
		})
	}
}

func (b *AbstractShell) writeCloneCmd(w ShellWriter, build *common.Build, projectDir string) {
	w.RmDir(projectDir)
	if depth := build.GetGitDepth(); depth != "" {
		w.Notice("Cloning repository for %s with git depth set to %s...", build.RefName, depth)
		w.Command("git", "clone", build.RepoURL, projectDir, "--depth", depth, "--branch", build.RefName)
	} else {
		w.Notice("Cloning repository...")
		w.Command("git", "clone", build.RepoURL, projectDir)
	}
	w.Cd(projectDir)
}

func (b *AbstractShell) writeFetchCmd(w ShellWriter, build *common.Build, projectDir string, gitDir string) {
	depth := build.GetGitDepth()

	w.IfDirectory(gitDir)
	if depth != "" {
		w.Notice("Fetching changes for %s with git depth set to %s...", build.RefName, depth)
	} else {
		w.Notice("Fetching changes...")
	}
	w.Cd(projectDir)
	w.Command("git", "clean", "-ffdx")
	w.Command("git", "reset", "--hard")
	w.Command("git", "remote", "set-url", "origin", build.RepoURL)
	if depth != "" {
		var refspec string
		if build.Tag {
			refspec = "+refs/tags/" + build.RefName + ":refs/tags/" + build.RefName
		} else {
			refspec = "+refs/heads/" + build.RefName + ":refs/remotes/origin/" + build.RefName
		}
		w.Command("git", "fetch", "--depth", depth, "origin", "--prune", refspec)
	} else {
		w.Command("git", "fetch", "origin", "--prune", "+refs/heads/*:refs/remotes/origin/*", "+refs/tags/*:refs/tags/*")
	}
	w.Else()
	b.writeCloneCmd(w, build, projectDir)
	w.EndIf()
}

func (b *AbstractShell) writeCheckoutCmd(w ShellWriter, build *common.Build) {
	w.Notice("Checking out %s as %s...", build.Sha[0:8], build.RefName)
	// We remove a git index file, this is required if `git checkout` is terminated
	w.RmFile(".git/index.lock")
	w.Command("git", "checkout", "-q", build.Sha)
}

func (b *AbstractShell) cacheFile(build *common.Build, userKey string) (key, file string) {
	if build.CacheDir == "" {
		return
	}

	// Deduce cache key
	key = path.Join(build.Name, build.RefName)
	if userKey != "" {
		key = build.GetAllVariables().ExpandValue(userKey)
	}

	// Ignore cache without the key
	if key == "" {
		return
	}

	file = path.Join(build.CacheDir, key, "cache.zip")
	file, err := filepath.Rel(build.BuildDir, file)
	if err != nil {
		return "", ""
	}
	return
}

func (o *archivingOptions) CommandArguments() (args []string) {
	for _, path := range o.Paths {
		args = append(args, "--path", path)
	}

	if o.Untracked {
		args = append(args, "--untracked")
	}
	return
}

func (b *AbstractShell) guardRunnerCommand(w ShellWriter, runnerCommand string, action string, f func()) {
	if runnerCommand == "" {
		w.Warning("%s is not supported by this executor.", action)
		return
	}

	w.IfCmd(runnerCommand, "--version")
	f()
	w.Else()
	w.Warning("Missing %s. %s is disabled.", runnerCommand, action)
	w.EndIf()
}

func (b *AbstractShell) cacheExtractor(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}

	// Skip restoring cache if no cache is defined
	if archiverArgs := options.CommandArguments(); len(archiverArgs) == 0 {
		return
	}

	// Skip archiving if no cache is defined
	cacheKey, cacheFile := b.cacheFile(info.Build, options.Key)
	if cacheKey == "" {
		return
	}

	args := []string{
		"cache-extractor",
		"--file", cacheFile,
	}

	// Generate cache download address
	if url := getCacheDownloadURL(info.Build, cacheKey); url != nil {
		args = append(args, "--url", url.String())
	}

	// Execute archive command
	b.guardRunnerCommand(w, info.RunnerCommand, "Extracting cache", func() {
		w.Notice("Checking cache for %s...", cacheKey)
		w.Command(info.RunnerCommand, args...)
	})
}

func (b *AbstractShell) downloadArtifacts(w ShellWriter, build *common.BuildInfo, info common.ShellScriptInfo) {
	args := []string{
		"artifacts-downloader",
		"--url",
		info.Build.Runner.URL,
		"--token",
		build.Token,
		"--id",
		strconv.Itoa(build.ID),
	}

	w.Notice("Downloading artifacts for %s (%d)...", build.Name, build.ID)
	w.Command(info.RunnerCommand, args...)
}

func (b *AbstractShell) buildArtifacts(dependencies *dependencies, info common.ShellScriptInfo) (otherBuilds []common.BuildInfo) {
	for _, otherBuild := range info.Build.DependsOnBuilds {
		if otherBuild.Artifacts == nil || otherBuild.Artifacts.Filename == "" {
			continue
		}
		if !dependencies.IsDependent(otherBuild.Name) {
			continue
		}
		otherBuilds = append(otherBuilds, otherBuild)
	}
	return
}

func (b *AbstractShell) downloadAllArtifacts(w ShellWriter, dependencies *dependencies, info common.ShellScriptInfo) {
	otherBuilds := b.buildArtifacts(dependencies, info)
	if len(otherBuilds) == 0 {
		return
	}

	b.guardRunnerCommand(w, info.RunnerCommand, "Artifacts downloading", func() {
		for _, otherBuild := range otherBuilds {
			b.downloadArtifacts(w, &otherBuild, info)
		}
	})
}

func (b *AbstractShell) writePrepareScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	b.writeExports(w, info)

	build := info.Build
	projectDir := build.FullProjectDir()
	gitDir := path.Join(build.FullProjectDir(), ".git")

	b.writeTLSCAInfo(w, info.Build, "GIT_SSL_CAINFO")
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	w.Command("git", "config", "--global", "fetch.recurseSubmodules", "false")
	switch info.Build.GetGitStrategy() {
	case common.GitFetch:
		b.writeFetchCmd(w, build, projectDir, gitDir)

	case common.GitClone:
		b.writeCloneCmd(w, build, projectDir)

	default:
		return errors.New("unknown GIT_STRATEGY")
	}

	b.writeCheckoutCmd(w, build)

	// Parse options
	var options shellOptions
	err = info.Build.Options.Decode(&options)
	if err != nil {
		return
	}

	// Try to restore from main cache, if not found cache for master
	b.cacheExtractor(w, options.Cache, info)

	// Process all artifacts
	b.downloadAllArtifacts(w, options.Dependencies, info)
	return nil
}

func (b *AbstractShell) writeBuildScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)

	commands := info.Build.Commands
	commands = strings.TrimSpace(commands)
	for _, command := range strings.Split(commands, "\n") {
		command = strings.TrimSpace(command)
		if command != "" {
			w.Notice("$ %s", command)
		} else {
			w.EmptyLine()
		}
		w.Line(command)
		w.CheckForErrors()
	}

	return nil
}

func (b *AbstractShell) cacheArchiver(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}

	// Skip archiving if no cache is defined
	cacheKey, cacheFile := b.cacheFile(info.Build, options.Key)
	if cacheKey == "" {
		return
	}

	args := []string{
		"cache-archiver",
		"--file", cacheFile,
	}

	// Create list of files to archive
	archiverArgs := options.CommandArguments()
	if len(archiverArgs) == 0 {
		// Skip creating archive
		return
	}
	args = append(args, archiverArgs...)

	// Generate cache upload address
	if url := getCacheUploadURL(info.Build, cacheKey); url != nil {
		args = append(args, "--url", url.String())
	}

	b.guardRunnerCommand(w, info.RunnerCommand, "Creating cache", func() {
		// Execute archive command
		w.Notice("Creating cache %s...", cacheKey)
		w.Command(info.RunnerCommand, args...)
	})
}

func (b *AbstractShell) uploadArtifacts(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}
	if info.Build.Runner.URL == "" {
		return
	}

	args := []string{
		"artifacts-uploader",
		"--url",
		info.Build.Runner.URL,
		"--token",
		info.Build.Token,
		"--id",
		strconv.Itoa(info.Build.ID),
	}

	// Create list of files to archive
	archiverArgs := options.CommandArguments()
	if len(archiverArgs) == 0 {
		// Skip creating archive
		return
	}
	args = append(args, archiverArgs...)

	// Get artifacts:name
	if name, ok := info.Build.Options.GetString("artifacts", "name"); ok && name != "" {
		args = append(args, "--name", name)
	}

	// Get artifacts:expire_in
	if expireIn, ok := info.Build.Options.GetString("artifacts", "expire_in"); ok && expireIn != "" {
		args = append(args, "--expire-in", expireIn)
	}

	b.guardRunnerCommand(w, info.RunnerCommand, "Uploading artifacts", func() {
		w.Notice("Uploading artifacts...")
		w.Command(info.RunnerCommand, args...)
	})
}

func (b *AbstractShell) writeAfterScript(w ShellWriter, info common.ShellScriptInfo) error {
	shellOptions := struct {
		AfterScript []string `json:"after_script"`
	}{}
	err := info.Build.Options.Decode(&shellOptions)
	if err != nil {
		return err
	}

	if len(shellOptions.AfterScript) == 0 {
		return nil
	}

	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)

	w.Notice("Running after script...")

	for _, command := range shellOptions.AfterScript {
		command = strings.TrimSpace(command)
		if command != "" {
			w.Notice("$ %s", command)
		} else {
			w.EmptyLine()
		}
		w.Line(command)
		w.CheckForErrors()
	}

	return nil
}

func (b *AbstractShell) writeArchiveCacheScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	// Parse options
	var options shellOptions
	err = info.Build.Options.Decode(&options)
	if err != nil {
		return
	}

	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	// Find cached files and archive them
	b.cacheArchiver(w, options.Cache, info)
	return
}

func (b *AbstractShell) writeUploadArtifactsScript(w ShellWriter, info common.ShellScriptInfo) (err error) {
	// Parse options
	var options shellOptions
	err = info.Build.Options.Decode(&options)
	if err != nil {
		return
	}

	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	// Upload artifacts
	b.uploadArtifacts(w, options.Artifacts, info)
	return
}

func (b *AbstractShell) writeScript(w ShellWriter, scriptType common.ShellScriptType, info common.ShellScriptInfo) (err error) {
	switch scriptType {
	case common.ShellPrepareScript:
		return b.writePrepareScript(w, info)

	case common.ShellBuildScript:
		return b.writeBuildScript(w, info)

	case common.ShellAfterScript:
		return b.writeAfterScript(w, info)

	case common.ShellArchiveCache:
		return b.writeArchiveCacheScript(w, info)

	case common.ShellUploadArtifacts:
		return b.writeUploadArtifactsScript(w, info)

	default:
		return errors.New("Not supported script type: " + string(scriptType))
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ssh

import (
	"crypto"
	"crypto/rand"
	"fmt"
	"io"
	"sync"

	_ "crypto/sha1"
	_ "crypto/sha256"
	_ "crypto/sha512"
)

// These are string constants in the SSH protocol.
const (
	compressionNone = "none"
	serviceUserAuth = "ssh-userauth"
	serviceSSH      = "ssh-connection"
)

// supportedCiphers specifies the supported ciphers in preference order.
var supportedCiphers = []string{
	"aes128-ctr", "aes192-ctr", "aes256-ctr",
	"aes128-gcm@openssh.com",
	"arcfour256", "arcfour128",
}

// supportedKexAlgos specifies the supported key-exchange algorithms in
// preference order.
var supportedKexAlgos = []string{
	kexAlgoCurve25519SHA256,
	// P384 and P521 are not constant-time yet, but since we don't
	// reuse ephemeral keys, using them for ECDH should be OK.
	kexAlgoECDH256, kexAlgoECDH384, kexAlgoECDH521,
	kexAlgoDH14SHA1, kexAlgoDH1SHA1,
}

// supportedKexAlgos specifies the supported host-key algorithms (i.e. methods
// of authenticating servers) in preference order.
var supportedHostKeyAlgos = []string{
	CertAlgoRSAv01, CertAlgoDSAv01, CertAlgoECDSA256v01,
	CertAlgoECDSA384v01, CertAlgoECDSA521v01, CertAlgoED25519v01,

	KeyAlgoECDSA256, KeyAlgoECDSA384, KeyAlgoECDSA521,
	KeyAlgoRSA, KeyAlgoDSA,

	KeyAlgoED25519,
}

// supportedMACs specifies a default set of MAC algorithms in preference order.
// This is based on RFC 4253, section 6.4, but with hmac-md5 variants removed
// because they have reached the end of their useful life.
var supportedMACs = []string{
	"hmac-sha2-256-etm@openssh.com", "hmac-sha2-256", "hmac-sha1", "hmac-sha1-96",
}

var supportedCompressions = []string{compressionNone}

// hashFuncs keeps the mapping of supported algorithms to their respective
// hashes needed for signature verification.
var hashFuncs = map[string]crypto.Hash{
	KeyAlgoRSA:          crypto.SHA1,
	KeyAlgoDSA:          crypto.SHA1,
	KeyAlgoECDSA256:     crypto.SHA256,
	KeyAlgoECDSA384:     crypto.SHA384,
	KeyAlgoECDSA521:     crypto.SHA512,
	CertAlgoRSAv01:      crypto.SHA1,
	CertAlgoDSAv01:      crypto.SHA1,
	CertAlgoECDSA256v01: crypto.SHA256,
	CertAlgoECDSA384v01: crypto.SHA384,
	CertAlgoECDSA521v01: crypto.SHA512,
}

// unexpectedMessageError results when the SSH message that we received didn't
// match what we wanted.
func unexpectedMessageError(expected, got uint8) error {
	return fmt.Errorf("ssh: unexpected message type %d (expected %d)", got, expected)
}

// parseError results from a malformed SSH message.
func parseError(tag uint8) error {
	return fmt.Errorf("ssh: parse error in message type %d", tag)
}

func findCommon(what string, client []string, server []string) (common string, err error) {
	for _, c := range client {
		for _, s := range server {
			if c == s {
				return c, nil
			}
		}
	}
	return "", fmt.Errorf("ssh: no common algorithm for %s; client offered: %v, server offered: %v", what, client, server)
}

type directionAlgorithms struct {
	Cipher      string
	MAC         string
	Compression string
}

// rekeyBytes returns a rekeying intervals in bytes.
func (a *directionAlgorithms) rekeyBytes() int64 {
	// According to RFC4344 block ciphers should rekey after
	// 2^(BLOCKSIZE/4) blocks. For all AES flavors BLOCKSIZE is
	// 128.
	switch a.Cipher {
	case "aes128-ctr", "aes192-ctr", "aes256-ctr", gcmCipherID, aes128cbcID:
		return 16 * (1 << 32)

	}

	// For others, stick with RFC4253 recommendation to rekey after 1 Gb of data.
	return 1 << 30
}

type algorithms struct {
	kex     string
	hostKey string
	w       directionAlgorithms
	r       directionAlgorithms
}

func findAgreedAlgorithms(clientKexInit, serverKexInit *kexInitMsg) (algs *algorithms, err error) {
	result := &algorithms{}

	result.kex, err = findCommon("key exchange", clientKexInit.KexAlgos, serverKexInit.KexAlgos)
	if err != nil {
		return
	}

	result.hostKey, err = findCommon("host key", clientKexInit.ServerHostKeyAlgos, serverKexInit.ServerHostKeyAlgos)
	if err != nil {
		return
	}

	result.w.Cipher, err = findCommon("client to server cipher", clientKexInit.CiphersClientServer, serverKexInit.CiphersClientServer)
	if err != nil {
		return
	}

	result.r.Cipher, err = findCommon("server to client cipher", clientKexInit.CiphersServerClient, serverKexInit.CiphersServerClient)
	if err != nil {
		return
	}

	result.w.MAC, err = findCommon("client to server MAC", clientKexInit.MACsClientServer, serverKexInit.MACsClientServer)
	if err != nil {
		return
	}

	result.r.MAC, err = findCommon("server to client MAC", clientKexInit.MACsServerClient, serverKexInit.MACsServerClient)
	if err != nil {
		return
	}

	result.w.Compression, err = findCommon("client to server compression", clientKexInit.CompressionClientServer, serverKexInit.CompressionClientServer)
	if err != nil {
		return
	}

	result.r.Compression, err = findCommon("server to client compression", clientKexInit.CompressionServerClient, serverKexInit.CompressionServerClient)
	if err != nil {
		return
	}

	return result, nil
}

// If rekeythreshold is too small, we can't make any progress sending
// stuff.
const minRekeyThreshold uint64 = 256

// Config contains configuration data common to both ServerConfig and
// ClientConfig.
type Config struct {
	// Rand provides the source of entropy for cryptographic
	// primitives. If Rand is nil, the cryptographic random reader
	// in package crypto/rand will be used.
	Rand io.Reader

	// The maximum number of bytes sent or received after which a
	// new key is negotiated. It must be at least 256. If
	// unspecified, 1 gigabyte is used.
	RekeyThreshold uint64

	// The allowed key exchanges algorithms. If unspecified then a
	// default set of algorithms is used.
	KeyExchanges []string

	// The allowed cipher algorithms. If unspecified then a sensible
	// default is used.
	Ciphers []string

	// The allowed MAC algorithms. If unspecified then a sensible default
	// is used.
	MACs []string
}

// SetDefaults sets sensible values for unset fields in config. This is
// exported for testing: Configs passed to SSH functions are copied and have
// default values set automatically.
func (c *Config) SetDefaults() {
	if c.Rand == nil {
		c.Rand = rand.Reader
	}
	if c.Ciphers == nil {
		c.Ciphers = supportedCiphers
	}
	var ciphers []string
	for _, c := range c.Ciphers {
		if cipherModes[c] != nil {
			// reject the cipher if we have no cipherModes definition
			ciphers = append(ciphers, c)
		}
	}
	c.Ciphers = ciphers

	if c.KeyExchanges == nil {
		c.KeyExchanges = supportedKexAlgos
	}

	if c.MACs == nil {
		c.MACs = supportedMACs
	}

	if c.RekeyThreshold == 0 {
		// RFC 4253, section 9 suggests rekeying after 1G.
		c.RekeyThreshold = 1 << 30
	}
	if c.RekeyThreshold < minRekeyThreshold {
		c.RekeyThreshold = minRekeyThreshold
	}
}

// buildDataSignedForAuth returns the data that is signed in order to prove
// possession of a private key. See RFC 4252, section 7.
func buildDataSignedForAuth(sessionId []byte, req userAuthRequestMsg, algo, pubKey []byte) []byte {
	data := struct {
		Session []byte
		Type    byte
		User    string
		Service string
		Method  string
		Sign    bool
		Algo    []byte
		PubKey  []byte
	}{
		sessionId,
		msgUserAuthRequest,
		req.User,
		req.Service,
		req.Method,
		true,
		algo,
		pubKey,
	}
	return Marshal(data)
}

func appendU16(buf []byte, n uint16) []byte {
	return append(buf, byte(n>>8), byte(n))
}

func appendU32(buf []byte, n uint32) []byte {
	return append(buf, byte(n>>24), byte(n>>16), byte(n>>8), byte(n))
}

func appendU64(buf []byte, n uint64) []byte {
	return append(buf,
		byte(n>>56), byte(n>>48), byte(n>>40), byte(n>>32),
		byte(n>>24), byte(n>>16), byte(n>>8), byte(n))
}

func appendInt(buf []byte, n int) []byte {
	return appendU32(buf, uint32(n))
}

func appendString(buf []byte, s string) []byte {
	buf = appendU32(buf, uint32(len(s)))
	buf = append(buf, s...)
	return buf
}

func appendBool(buf []byte, b bool) []byte {
	if b {
		return append(buf, 1)
	}
	return append(buf, 0)
}

// newCond is a helper to hide the fact that there is no usable zero
// value for sync.Cond.
func newCond() *sync.Cond { return sync.NewCond(new(sync.Mutex)) }

// window represents the buffer available to clients
// wishing to write to a channel.
type window struct {
	*sync.Cond
	win          uint32 // RFC 4254 5.2 says the window size can grow to 2^32-1
	writeWaiters int
	closed       bool
}

// add adds win to the amount of window available
// for consumers.
func (w *window) add(win uint32) bool {
	// a zero sized window adjust is a noop.
	if win == 0 {
		return true
	}
	w.L.Lock()
	if w.win+win < win {
		w.L.Unlock()
		return false
	}
	w.win += win
	// It is unusual that multiple goroutines would be attempting to reserve
	// window space, but not guaranteed. Use broadcast to notify all waiters
	// that additional window is available.
	w.Broadcast()
	w.L.Unlock()
	return true
}

// close sets the window to closed, so all reservations fail
// immediately.
func (w *window) close() {
	w.L.Lock()
	w.closed = true
	w.Broadcast()
	w.L.Unlock()
}

// reserve reserves win from the available window capacity.
// If no capacity remains, reserve will block. reserve may
// return less than requested.
func (w *window) reserve(win uint32) (uint32, error) {
	var err error
	w.L.Lock()
	w.writeWaiters++
	w.Broadcast()
	for w.win == 0 && !w.closed {
		w.Wait()
	}
	w.writeWaiters--
	if w.win < win {
		win = w.win
	}
	w.win -= win
	if w.closed {
		err = io.EOF
	}
	w.L.Unlock()
	return win, err
}

// waitWriterBlocked waits until some goroutine is blocked for further
// writes. It is used in tests only.
func (w *window) waitWriterBlocked() {
	w.Cond.L.Lock()
	for w.writeWaiters == 0 {
		w.Cond.Wait()
	}
	w.Cond.L.Unlock()
}
<commit_msg>ssh: fix typo in unexported comment<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ssh

import (
	"crypto"
	"crypto/rand"
	"fmt"
	"io"
	"sync"

	_ "crypto/sha1"
	_ "crypto/sha256"
	_ "crypto/sha512"
)

// These are string constants in the SSH protocol.
const (
	compressionNone = "none"
	serviceUserAuth = "ssh-userauth"
	serviceSSH      = "ssh-connection"
)

// supportedCiphers specifies the supported ciphers in preference order.
var supportedCiphers = []string{
	"aes128-ctr", "aes192-ctr", "aes256-ctr",
	"aes128-gcm@openssh.com",
	"arcfour256", "arcfour128",
}

// supportedKexAlgos specifies the supported key-exchange algorithms in
// preference order.
var supportedKexAlgos = []string{
	kexAlgoCurve25519SHA256,
	// P384 and P521 are not constant-time yet, but since we don't
	// reuse ephemeral keys, using them for ECDH should be OK.
	kexAlgoECDH256, kexAlgoECDH384, kexAlgoECDH521,
	kexAlgoDH14SHA1, kexAlgoDH1SHA1,
}

// supportedHostKeyAlgos specifies the supported host-key algorithms (i.e. methods
// of authenticating servers) in preference order.
var supportedHostKeyAlgos = []string{
	CertAlgoRSAv01, CertAlgoDSAv01, CertAlgoECDSA256v01,
	CertAlgoECDSA384v01, CertAlgoECDSA521v01, CertAlgoED25519v01,

	KeyAlgoECDSA256, KeyAlgoECDSA384, KeyAlgoECDSA521,
	KeyAlgoRSA, KeyAlgoDSA,

	KeyAlgoED25519,
}

// supportedMACs specifies a default set of MAC algorithms in preference order.
// This is based on RFC 4253, section 6.4, but with hmac-md5 variants removed
// because they have reached the end of their useful life.
var supportedMACs = []string{
	"hmac-sha2-256-etm@openssh.com", "hmac-sha2-256", "hmac-sha1", "hmac-sha1-96",
}

var supportedCompressions = []string{compressionNone}

// hashFuncs keeps the mapping of supported algorithms to their respective
// hashes needed for signature verification.
var hashFuncs = map[string]crypto.Hash{
	KeyAlgoRSA:          crypto.SHA1,
	KeyAlgoDSA:          crypto.SHA1,
	KeyAlgoECDSA256:     crypto.SHA256,
	KeyAlgoECDSA384:     crypto.SHA384,
	KeyAlgoECDSA521:     crypto.SHA512,
	CertAlgoRSAv01:      crypto.SHA1,
	CertAlgoDSAv01:      crypto.SHA1,
	CertAlgoECDSA256v01: crypto.SHA256,
	CertAlgoECDSA384v01: crypto.SHA384,
	CertAlgoECDSA521v01: crypto.SHA512,
}

// unexpectedMessageError results when the SSH message that we received didn't
// match what we wanted.
func unexpectedMessageError(expected, got uint8) error {
	return fmt.Errorf("ssh: unexpected message type %d (expected %d)", got, expected)
}

// parseError results from a malformed SSH message.
func parseError(tag uint8) error {
	return fmt.Errorf("ssh: parse error in message type %d", tag)
}

func findCommon(what string, client []string, server []string) (common string, err error) {
	for _, c := range client {
		for _, s := range server {
			if c == s {
				return c, nil
			}
		}
	}
	return "", fmt.Errorf("ssh: no common algorithm for %s; client offered: %v, server offered: %v", what, client, server)
}

type directionAlgorithms struct {
	Cipher      string
	MAC         string
	Compression string
}

// rekeyBytes returns a rekeying intervals in bytes.
func (a *directionAlgorithms) rekeyBytes() int64 {
	// According to RFC4344 block ciphers should rekey after
	// 2^(BLOCKSIZE/4) blocks. For all AES flavors BLOCKSIZE is
	// 128.
	switch a.Cipher {
	case "aes128-ctr", "aes192-ctr", "aes256-ctr", gcmCipherID, aes128cbcID:
		return 16 * (1 << 32)

	}

	// For others, stick with RFC4253 recommendation to rekey after 1 Gb of data.
	return 1 << 30
}

type algorithms struct {
	kex     string
	hostKey string
	w       directionAlgorithms
	r       directionAlgorithms
}

func findAgreedAlgorithms(clientKexInit, serverKexInit *kexInitMsg) (algs *algorithms, err error) {
	result := &algorithms{}

	result.kex, err = findCommon("key exchange", clientKexInit.KexAlgos, serverKexInit.KexAlgos)
	if err != nil {
		return
	}

	result.hostKey, err = findCommon("host key", clientKexInit.ServerHostKeyAlgos, serverKexInit.ServerHostKeyAlgos)
	if err != nil {
		return
	}

	result.w.Cipher, err = findCommon("client to server cipher", clientKexInit.CiphersClientServer, serverKexInit.CiphersClientServer)
	if err != nil {
		return
	}

	result.r.Cipher, err = findCommon("server to client cipher", clientKexInit.CiphersServerClient, serverKexInit.CiphersServerClient)
	if err != nil {
		return
	}

	result.w.MAC, err = findCommon("client to server MAC", clientKexInit.MACsClientServer, serverKexInit.MACsClientServer)
	if err != nil {
		return
	}

	result.r.MAC, err = findCommon("server to client MAC", clientKexInit.MACsServerClient, serverKexInit.MACsServerClient)
	if err != nil {
		return
	}

	result.w.Compression, err = findCommon("client to server compression", clientKexInit.CompressionClientServer, serverKexInit.CompressionClientServer)
	if err != nil {
		return
	}

	result.r.Compression, err = findCommon("server to client compression", clientKexInit.CompressionServerClient, serverKexInit.CompressionServerClient)
	if err != nil {
		return
	}

	return result, nil
}

// If rekeythreshold is too small, we can't make any progress sending
// stuff.
const minRekeyThreshold uint64 = 256

// Config contains configuration data common to both ServerConfig and
// ClientConfig.
type Config struct {
	// Rand provides the source of entropy for cryptographic
	// primitives. If Rand is nil, the cryptographic random reader
	// in package crypto/rand will be used.
	Rand io.Reader

	// The maximum number of bytes sent or received after which a
	// new key is negotiated. It must be at least 256. If
	// unspecified, 1 gigabyte is used.
	RekeyThreshold uint64

	// The allowed key exchanges algorithms. If unspecified then a
	// default set of algorithms is used.
	KeyExchanges []string

	// The allowed cipher algorithms. If unspecified then a sensible
	// default is used.
	Ciphers []string

	// The allowed MAC algorithms. If unspecified then a sensible default
	// is used.
	MACs []string
}

// SetDefaults sets sensible values for unset fields in config. This is
// exported for testing: Configs passed to SSH functions are copied and have
// default values set automatically.
func (c *Config) SetDefaults() {
	if c.Rand == nil {
		c.Rand = rand.Reader
	}
	if c.Ciphers == nil {
		c.Ciphers = supportedCiphers
	}
	var ciphers []string
	for _, c := range c.Ciphers {
		if cipherModes[c] != nil {
			// reject the cipher if we have no cipherModes definition
			ciphers = append(ciphers, c)
		}
	}
	c.Ciphers = ciphers

	if c.KeyExchanges == nil {
		c.KeyExchanges = supportedKexAlgos
	}

	if c.MACs == nil {
		c.MACs = supportedMACs
	}

	if c.RekeyThreshold == 0 {
		// RFC 4253, section 9 suggests rekeying after 1G.
		c.RekeyThreshold = 1 << 30
	}
	if c.RekeyThreshold < minRekeyThreshold {
		c.RekeyThreshold = minRekeyThreshold
	}
}

// buildDataSignedForAuth returns the data that is signed in order to prove
// possession of a private key. See RFC 4252, section 7.
func buildDataSignedForAuth(sessionId []byte, req userAuthRequestMsg, algo, pubKey []byte) []byte {
	data := struct {
		Session []byte
		Type    byte
		User    string
		Service string
		Method  string
		Sign    bool
		Algo    []byte
		PubKey  []byte
	}{
		sessionId,
		msgUserAuthRequest,
		req.User,
		req.Service,
		req.Method,
		true,
		algo,
		pubKey,
	}
	return Marshal(data)
}

func appendU16(buf []byte, n uint16) []byte {
	return append(buf, byte(n>>8), byte(n))
}

func appendU32(buf []byte, n uint32) []byte {
	return append(buf, byte(n>>24), byte(n>>16), byte(n>>8), byte(n))
}

func appendU64(buf []byte, n uint64) []byte {
	return append(buf,
		byte(n>>56), byte(n>>48), byte(n>>40), byte(n>>32),
		byte(n>>24), byte(n>>16), byte(n>>8), byte(n))
}

func appendInt(buf []byte, n int) []byte {
	return appendU32(buf, uint32(n))
}

func appendString(buf []byte, s string) []byte {
	buf = appendU32(buf, uint32(len(s)))
	buf = append(buf, s...)
	return buf
}

func appendBool(buf []byte, b bool) []byte {
	if b {
		return append(buf, 1)
	}
	return append(buf, 0)
}

// newCond is a helper to hide the fact that there is no usable zero
// value for sync.Cond.
func newCond() *sync.Cond { return sync.NewCond(new(sync.Mutex)) }

// window represents the buffer available to clients
// wishing to write to a channel.
type window struct {
	*sync.Cond
	win          uint32 // RFC 4254 5.2 says the window size can grow to 2^32-1
	writeWaiters int
	closed       bool
}

// add adds win to the amount of window available
// for consumers.
func (w *window) add(win uint32) bool {
	// a zero sized window adjust is a noop.
	if win == 0 {
		return true
	}
	w.L.Lock()
	if w.win+win < win {
		w.L.Unlock()
		return false
	}
	w.win += win
	// It is unusual that multiple goroutines would be attempting to reserve
	// window space, but not guaranteed. Use broadcast to notify all waiters
	// that additional window is available.
	w.Broadcast()
	w.L.Unlock()
	return true
}

// close sets the window to closed, so all reservations fail
// immediately.
func (w *window) close() {
	w.L.Lock()
	w.closed = true
	w.Broadcast()
	w.L.Unlock()
}

// reserve reserves win from the available window capacity.
// If no capacity remains, reserve will block. reserve may
// return less than requested.
func (w *window) reserve(win uint32) (uint32, error) {
	var err error
	w.L.Lock()
	w.writeWaiters++
	w.Broadcast()
	for w.win == 0 && !w.closed {
		w.Wait()
	}
	w.writeWaiters--
	if w.win < win {
		win = w.win
	}
	w.win -= win
	if w.closed {
		err = io.EOF
	}
	w.L.Unlock()
	return win, err
}

// waitWriterBlocked waits until some goroutine is blocked for further
// writes. It is used in tests only.
func (w *window) waitWriterBlocked() {
	w.Cond.L.Lock()
	for w.writeWaiters == 0 {
		w.Cond.Wait()
	}
	w.Cond.L.Unlock()
}
<|endoftext|>


<commit_before>package ssh

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/blacknon/go-sshlib"
	"github.com/c-bata/go-prompt"
	"github.com/c-bata/go-prompt/completer"
)

// Pshell is Parallel-Shell struct
type pShell struct {
	Signal      chan os.Signal
	Count       int
	ServerList  []string
	Connects    []*psConnect
	PROMPT      string
	History     map[int]map[string]*pShellHistory
	HistoryFile string
	Complete    []prompt.Suggest
}

// psConnect is pShell connect struct.
type psConnect struct {
	Name   string
	Output *Output
	*sshlib.Connect
}

// variable
var (
	defaultPrompt      = "[${COUNT}] <<< "          // Default PROMPT
	defaultOPrompt     = "[${SERVER}][${COUNT}] > " // Default OPROMPT
	defaultHistoryFile = "~/.lssh_history"          // Default Parallel shell history file
)

func (r *Run) pshell() (err error) {
	// print header
	fmt.Println("Start parallel-shell...")
	r.printSelectServer()

	// read shell config
	config := r.Conf.Shell

	// overwrite default value config.Prompt
	if config.Prompt == "" {
		config.Prompt = defaultPrompt
	}

	// overwrite default value config.OPrompt
	if config.OPrompt == "" {
		config.OPrompt = defaultOPrompt
	}

	// overwrite default parallel shell history file
	if config.HistoryFile == "" {
		config.HistoryFile = defaultHistoryFile
	}

	// run pre cmd
	runCmdLocal(config.PreCmd)
	defer runCmdLocal(config.PostCmd)

	// Connect
	var cons []*psConnect
	for _, server := range r.ServerList {
		// Create *sshlib.Connect
		con, err := r.createSshConnect(server)
		if err != nil {
			log.Println(err)
			continue
		}

		// Create Output
		o := &Output{
			Templete:   config.OPrompt,
			ServerList: r.ServerList,
			Conf:       r.Conf.Server[server],
			AutoColor:  true,
		}

		// Create output prompt
		o.Create(server)

		psCon := &psConnect{
			Name:    server,
			Output:  o,
			Connect: con,
		}
		cons = append(cons, psCon)
	}

	// count sshlib.Connect.
	if len(cons) == 0 {
		return
	}

	// create new shell struct
	ps := &pShell{
		Signal:      make(chan os.Signal),
		ServerList:  r.ServerList,
		Connects:    cons,
		PROMPT:      config.Prompt,
		History:     map[int]map[string]*pShellHistory{},
		HistoryFile: config.HistoryFile,
	}

	// set signal
	signal.Notify(ps.Signal, syscall.SIGTERM, syscall.SIGINT)

	// old history list
	var historyCommand []string
	oldHistory, err := ps.GetHistoryFromFile()
	if err == nil {
		for _, h := range oldHistory {
			historyCommand = append(historyCommand, h.Command)
		}
	}

	// create complete data
	ps.GetCompleteData()

	// create prompt
	pShellPrompt, _ := ps.CreatePrompt()

	// create go-prompt
	p := prompt.New(
		ps.Executor,
		ps.Completer,
		prompt.OptionHistory(historyCommand),
		prompt.OptionPrefix(pShellPrompt),
		prompt.OptionLivePrefix(ps.CreatePrompt),
		prompt.OptionInputTextColor(prompt.Green),
		prompt.OptionPrefixTextColor(prompt.Blue),
		prompt.OptionCompletionWordSeparator(completer.FilePathCompletionSeparator), // test
	)

	// start go-prompt
	p.Run()

	return
}

// CreatePrompt is create shell prompt.
// default value is `[${COUNT}] <<< `
func (ps *pShell) CreatePrompt() (p string, result bool) {
	// set prompt templete (from conf)
	p = ps.PROMPT
	if p == "" {
		p = defaultPrompt
	}

	// Get env
	hostname, _ := os.Hostname()
	username := os.Getenv("USER")
	pwd := os.Getenv("PWD")

	// replace variable value
	p = strings.Replace(p, "${COUNT}", strconv.Itoa(ps.Count), -1)
	p = strings.Replace(p, "${HOSTNAME}", hostname, -1)
	p = strings.Replace(p, "${USER}", username, -1)
	p = strings.Replace(p, "${PWD}", pwd, -1)

	return p, true
}

// Executor run ssh command in parallel-shell.
func (ps *pShell) Executor(command string) {
	// trim space
	command = strings.TrimSpace(command)

	// local command regex
	localCmdRegex_out := regexp.MustCompile(`^%out.*$`)

	// check local command
	switch {
	// only enter(Ctrl + M)
	case command == "":
		return

	// exit or quit
	case command == "exit", command == "quit":
		// put history and exit
		ps.PutHistoryFile(command)
		os.Exit(0)

	// clear
	case command == "clear":
		fmt.Printf("\033[H\033[2J")
		ps.PutHistoryFile(command)
		return

	// history
	case command == "history":
		ps.localCmd_history()
		return

	// %out [num]
	case localCmdRegex_out.MatchString(command):
		cmdSlice := strings.SplitN(command, " ", 2)

		// set default num
		num := 0
		if ps.Count > 0 {
			num = ps.Count - 1
		}

		// get args
		if len(cmdSlice) > 1 {
			inum, err := strconv.Atoi(cmdSlice[1])
			if err != nil {
				return
			}

			// if num > count
			if inum >= ps.Count {
				return
			}

			num = inum
		}

		ps.localCmd_out(num)
		fmt.Println()
		return
	}

	// put history
	ps.PutHistoryFile(command)

	// run command
	ps.Run(command)

	return
}

// Completer lssh-shell complete function
func (ps *pShell) Completer(t prompt.Document) []prompt.Suggest {
	// TODO(blacknon): とりあえず値を仮置き。後で以下の処理を追加する(優先度A)
	//        - compgen(confで補完用の結果を取得するためのコマンドは指定可能にする)での補完結果の定期取得処理(+補完の取得用ローカルコマンドの追加)
	//        - compgenの結果をStructに保持させる
	//        - Structに保持されている補完内容をベースにCompleteの結果を返す
	//        - 何も入力していない場合は非表示にさせたい
	//        - ファイルについても対応させたい
	//        - ファイルやコマンドなど、状況に応じて補完対象を変えるにはやはり構文解析が必要になってくる。Parserを実装するまではコマンドのみ対応。
	//        	参考: https://github.com/c-bata/kube-prompt/blob/2276d167e2e693164c5980427a6809058a235c95/kube/completer.go

	// local command suggest
	localCmdSuggest := []prompt.Suggest{
		{Text: "exit", Description: "exit lssh shell"},
		{Text: "quit", Description: "exit lssh shell"},
		{Text: "clear", Description: "clear screen"},
		{Text: "history", Description: "show history"},
		{Text: "%out", Description: "%out [num], show history result."},
		// {Text: "%outlist", Description: "%outlist, show history result list."}, // outのリストを出力ためのローカルコマンド

		// outの出力でdiffをするためのローカルコマンド。すべての出力と比較するのはあまりに辛いと思われるため、最初の出力との比較、といった方式で対応するのが良いか？？
		// {Text: "%diff", Description: "%diff [num], show history result list."},

		// outの出力でユニークな出力だけを表示するコマンド
		// {Text: "%unique", Description: "%unique [num], show history result list."},

		// outの出力で重複している出力だけを表示するコマンド
		// {Text: "%duplicate", Description: "%duplicate [num], show history result list."},
	}

	// get complete data
	c := ps.Complete
	c = append(c, localCmdSuggest...)

	return prompt.FilterHasPrefix(c, t.GetWordBeforeCursor(), false)
}

// GetCompleteData get command list remote machine.
func (ps *pShell) GetCompleteData() {
	// TODO(blacknon):
	//   - 構文解析して、ファイルの補完処理も行わせる
	//     - 引数にコマンドorファイルの種別を渡すようにする
	//   - 補完コマンドをconfigでオプションとして指定できるようにする
	//     - あまり無いだろうけど、bash以外をリモートで使ってる場合など(ashとかzsh(レア)など)

	// bash complete command. use `compgen`.
	compCmd := []string{"compgen", "-c"}
	command := strings.Join(compCmd, " ")

	// command map
	cmdMap := map[string][]string{}

	// append command to cmdMap
	for _, c := range ps.Connects {
		// Create buffer
		buf := new(bytes.Buffer)

		// Create session, and output to buffer
		session, _ := c.CreateSession()
		session.Stdout = buf

		// Run get complete command
		session.Run(command)

		// Scan and put completed command to map.
		sc := bufio.NewScanner(buf)
		for sc.Scan() {
			cmdMap[sc.Text()] = append(cmdMap[sc.Text()], c.Name)
		}
	}

	// cmdMap to suggest
	for cmd, hosts := range cmdMap {
		// join hosts
		h := strings.Join(hosts, ",")

		// create suggest
		suggest := prompt.Suggest{
			Text:        cmd,
			Description: "Command. from:" + h,
		}

		// append ps.Complete
		ps.Complete = append(ps.Complete, suggest)
	}
}

// Run is exec command.
func (ps *pShell) Run(command string) {
	// Create History
	ps.History[ps.Count] = map[string]*pShellHistory{}

	// create chanel
	finished := make(chan bool)    // Run Command finish channel
	input := make(chan io.Writer)  // Get io.Writer at input channel
	exitInput := make(chan bool)   // Input finish channel
	exitSignal := make(chan bool)  // Send kill signal finish channel
	exitHistory := make(chan bool) // Put History finish channel

	// create []io.Writer after in MultiWriter
	var writers []io.Writer

	// for connect and run
	for _, fc := range ps.Connects {
		// set variable c
		// NOTE: Variables need to be assigned separately for processing by goroutine.
		c := fc

		// Get output data channel
		output := make(chan []byte)
		// defer close(output)

		// Set count num
		c.Output.Count = ps.Count

		// Create output buffer, and MultiWriter
		buf := new(bytes.Buffer)
		omw := io.MultiWriter(os.Stdout, buf)
		c.Output.OutputWriter = omw

		// put result
		go ps.PutHistoryResult(c.Name, command, buf, exitHistory)

		// Run command
		go func() {
			c.CmdWriter(command, output, input)
			finished <- true
		}()

		// Get input(io.Writer), add MultiWriter
		w := <-input
		writers = append(writers, w)

		// run print Output
		go func() {
			printOutput(c.Output, output)
		}()
	}

	// create and run input writer
	mw := io.MultiWriter(writers...)
	go pushInput(exitInput, mw)

	// send kill signal function
	go ps.pushKillSignal(exitSignal, ps.Connects)

	// wait finished channel
	for i := 0; i < len(ps.Connects); i++ {
		<-finished
	}

	// Exit check signal
	exitSignal <- true

	// wait time (0.500 sec)
	time.Sleep(100 * time.Millisecond)

	// Exit Messages
	// Because it is Blocking.IO, you can not finish Input without input from the user.
	fmt.Fprintf(os.Stderr, "\n---\n%s\n", "Command exit. Please input Enter.")

	// Exit input
	exitInput <- true

	// Exit history
	m := new(sync.Mutex)
	for i := 0; i < len(ps.Connects); i++ {
		m.Lock()
		exitHistory <- true
		m.Unlock()
	}

	// Add count
	ps.Count += 1
}

// pushSignal is send kill signal to session.
func (ps *pShell) pushKillSignal(exitSig chan bool, conns []*psConnect) (err error) {
	i := 0
	for {
		select {
		case <-ps.Signal:
			if i == 0 {
				for _, c := range conns {
					// send kill
					c.Kill()
				}
				i = 1
			}
		case <-exitSig:
			return
		}
	}
	return
}
<commit_msg>bugfix<commit_after>package ssh

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/blacknon/go-sshlib"
	"github.com/c-bata/go-prompt"
	"github.com/c-bata/go-prompt/completer"
)

// Pshell is Parallel-Shell struct
type pShell struct {
	Signal      chan os.Signal
	Count       int
	ServerList  []string
	Connects    []*psConnect
	PROMPT      string
	History     map[int]map[string]*pShellHistory
	HistoryFile string
	Complete    []prompt.Suggest
}

// psConnect is pShell connect struct.
type psConnect struct {
	Name   string
	Output *Output
	*sshlib.Connect
}

// variable
var (
	defaultPrompt      = "[${COUNT}] <<< "          // Default PROMPT
	defaultOPrompt     = "[${SERVER}][${COUNT}] > " // Default OPROMPT
	defaultHistoryFile = "~/.lssh_history"          // Default Parallel shell history file
)

func (r *Run) pshell() (err error) {
	// print header
	fmt.Println("Start parallel-shell...")
	r.printSelectServer()

	// read shell config
	config := r.Conf.Shell

	// overwrite default value config.Prompt
	if config.Prompt == "" {
		config.Prompt = defaultPrompt
	}

	// overwrite default value config.OPrompt
	if config.OPrompt == "" {
		config.OPrompt = defaultOPrompt
	}

	// overwrite default parallel shell history file
	if config.HistoryFile == "" {
		config.HistoryFile = defaultHistoryFile
	}

	// run pre cmd
	runCmdLocal(config.PreCmd)
	defer runCmdLocal(config.PostCmd)

	// Connect
	var cons []*psConnect
	for _, server := range r.ServerList {
		// Create *sshlib.Connect
		con, err := r.createSshConnect(server)
		if err != nil {
			log.Println(err)
			continue
		}

		// Create Output
		o := &Output{
			Templete:   config.OPrompt,
			ServerList: r.ServerList,
			Conf:       r.Conf.Server[server],
			AutoColor:  true,
		}

		// Create output prompt
		o.Create(server)

		psCon := &psConnect{
			Name:    server,
			Output:  o,
			Connect: con,
		}
		cons = append(cons, psCon)
	}

	// count sshlib.Connect.
	if len(cons) == 0 {
		return
	}

	// create new shell struct
	ps := &pShell{
		Signal:      make(chan os.Signal),
		ServerList:  r.ServerList,
		Connects:    cons,
		PROMPT:      config.Prompt,
		History:     map[int]map[string]*pShellHistory{},
		HistoryFile: config.HistoryFile,
	}

	// set signal
	signal.Notify(ps.Signal, syscall.SIGTERM, syscall.SIGINT)

	// old history list
	var historyCommand []string
	oldHistory, err := ps.GetHistoryFromFile()
	if err == nil {
		for _, h := range oldHistory {
			historyCommand = append(historyCommand, h.Command)
		}
	}

	// create complete data
	ps.GetCompleteData()

	// create prompt
	pShellPrompt, _ := ps.CreatePrompt()

	// create go-prompt
	p := prompt.New(
		ps.Executor,
		ps.Completer,
		prompt.OptionHistory(historyCommand),
		prompt.OptionPrefix(pShellPrompt),
		prompt.OptionLivePrefix(ps.CreatePrompt),
		prompt.OptionInputTextColor(prompt.Green),
		prompt.OptionPrefixTextColor(prompt.Blue),
		prompt.OptionCompletionWordSeparator(completer.FilePathCompletionSeparator), // test
	)

	// start go-prompt
	p.Run()

	return
}

// CreatePrompt is create shell prompt.
// default value is `[${COUNT}] <<< `
func (ps *pShell) CreatePrompt() (p string, result bool) {
	// set prompt templete (from conf)
	p = ps.PROMPT
	if p == "" {
		p = defaultPrompt
	}

	// Get env
	hostname, _ := os.Hostname()
	username := os.Getenv("USER")
	pwd := os.Getenv("PWD")

	// replace variable value
	p = strings.Replace(p, "${COUNT}", strconv.Itoa(ps.Count), -1)
	p = strings.Replace(p, "${HOSTNAME}", hostname, -1)
	p = strings.Replace(p, "${USER}", username, -1)
	p = strings.Replace(p, "${PWD}", pwd, -1)

	return p, true
}

// Executor run ssh command in parallel-shell.
func (ps *pShell) Executor(command string) {
	// trim space
	command = strings.TrimSpace(command)

	// local command regex
	localCmdRegex_out := regexp.MustCompile(`^%out.*$`)

	// check local command
	switch {
	// only enter(Ctrl + M)
	case command == "":
		return

	// exit or quit
	case command == "exit", command == "quit":
		// put history and exit
		ps.PutHistoryFile(command)
		os.Exit(0)

	// clear
	case command == "clear":
		fmt.Printf("\033[H\033[2J")
		ps.PutHistoryFile(command)
		return

	// history
	case command == "history":
		ps.localCmd_history()
		return

	// %out [num]
	case localCmdRegex_out.MatchString(command):
		cmdSlice := strings.SplitN(command, " ", 2)

		// set default num
		num := 0
		if ps.Count > 0 {
			num = ps.Count - 1
		}

		// get args
		if len(cmdSlice) > 1 {
			inum, err := strconv.Atoi(cmdSlice[1])
			if err != nil {
				return
			}

			// if num > count
			if inum >= ps.Count {
				return
			}

			num = inum
		}

		ps.localCmd_out(num)
		fmt.Println()
		return
	}

	// put history
	ps.PutHistoryFile(command)

	// run command
	ps.Run(command)

	return
}

// Completer lssh-shell complete function
func (ps *pShell) Completer(t prompt.Document) []prompt.Suggest {
	// TODO(blacknon): とりあえず値を仮置き。後で以下の処理を追加する(優先度A)
	//        - compgen(confで補完用の結果を取得するためのコマンドは指定可能にする)での補完結果の定期取得処理(+補完の取得用ローカルコマンドの追加)
	//        - compgenの結果をStructに保持させる
	//        - Structに保持されている補完内容をベースにCompleteの結果を返す
	//        - 何も入力していない場合は非表示にさせたい
	//        - ファイルについても対応させたい
	//        - ファイルやコマンドなど、状況に応じて補完対象を変えるにはやはり構文解析が必要になってくる。Parserを実装するまではコマンドのみ対応。
	//        	参考: https://github.com/c-bata/kube-prompt/blob/2276d167e2e693164c5980427a6809058a235c95/kube/completer.go

	// local command suggest
	localCmdSuggest := []prompt.Suggest{
		{Text: "exit", Description: "exit lssh shell"},
		{Text: "quit", Description: "exit lssh shell"},
		{Text: "clear", Description: "clear screen"},
		{Text: "history", Description: "show history"},
		{Text: "%out", Description: "%out [num], show history result."},
		// {Text: "%outlist", Description: "%outlist, show history result list."}, // outのリストを出力ためのローカルコマンド

		// outの出力でdiffをするためのローカルコマンド。すべての出力と比較するのはあまりに辛いと思われるため、最初の出力との比較、といった方式で対応するのが良いか？？
		// {Text: "%diff", Description: "%diff [num], show history result list."},

		// outの出力でユニークな出力だけを表示するコマンド
		// {Text: "%unique", Description: "%unique [num], show history result list."},

		// outの出力で重複している出力だけを表示するコマンド
		// {Text: "%duplicate", Description: "%duplicate [num], show history result list."},
	}

	// get complete data
	c := ps.Complete
	c = append(c, localCmdSuggest...)

	return prompt.FilterHasPrefix(c, t.GetWordBeforeCursor(), false)
}

// GetCompleteData get command list remote machine.
func (ps *pShell) GetCompleteData() {
	// TODO(blacknon):
	//   - 構文解析して、ファイルの補完処理も行わせる
	//     - 引数にコマンドorファイルの種別を渡すようにする
	//   - 補完コマンドをconfigでオプションとして指定できるようにする
	//     - あまり無いだろうけど、bash以外をリモートで使ってる場合など(ashとかzsh(レア)など)

	// bash complete command. use `compgen`.
	compCmd := []string{"compgen", "-c"}
	command := strings.Join(compCmd, " ")

	// command map
	cmdMap := map[string][]string{}

	// append command to cmdMap
	for _, c := range ps.Connects {
		// Create buffer
		buf := new(bytes.Buffer)

		// Create session, and output to buffer
		session, _ := c.CreateSession()
		session.Stdout = buf

		// Run get complete command
		session.Run(command)

		// Scan and put completed command to map.
		sc := bufio.NewScanner(buf)
		for sc.Scan() {
			cmdMap[sc.Text()] = append(cmdMap[sc.Text()], c.Name)
		}
	}

	// cmdMap to suggest
	for cmd, hosts := range cmdMap {
		// join hosts
		h := strings.Join(hosts, ",")

		// create suggest
		suggest := prompt.Suggest{
			Text:        cmd,
			Description: "Command. from:" + h,
		}

		// append ps.Complete
		ps.Complete = append(ps.Complete, suggest)
	}
}

// Run is exec command.
func (ps *pShell) Run(command string) {
	// Create History
	ps.History[ps.Count] = map[string]*pShellHistory{}

	// create chanel
	finished := make(chan bool)    // Run Command finish channel
	input := make(chan io.Writer)  // Get io.Writer at input channel
	exitInput := make(chan bool)   // Input finish channel
	exitSignal := make(chan bool)  // Send kill signal finish channel
	exitHistory := make(chan bool) // Put History finish channel

	// create []io.Writer after in MultiWriter
	var writers []io.Writer

	// for connect and run
	m := new(sync.Mutex)
	for _, fc := range ps.Connects {

		// set variable c
		// NOTE: Variables need to be assigned separately for processing by goroutine.
		c := fc

		// Get output data channel
		output := make(chan []byte)
		// defer close(output)

		// Set count num
		c.Output.Count = ps.Count

		// Create output buffer, and MultiWriter
		buf := new(bytes.Buffer)
		omw := io.MultiWriter(os.Stdout, buf)
		c.Output.OutputWriter = omw

		// put result
		go func() {
			m.Lock()
			ps.PutHistoryResult(c.Name, command, buf, exitHistory)
			m.Unlock()
		}()

		// Run command
		go func() {
			c.CmdWriter(command, output, input)
			finished <- true
		}()

		// Get input(io.Writer), add MultiWriter
		w := <-input
		writers = append(writers, w)

		// run print Output
		go func() {
			printOutput(c.Output, output)
		}()
	}

	// create and run input writer
	mw := io.MultiWriter(writers...)
	go pushInput(exitInput, mw)

	// send kill signal function
	go ps.pushKillSignal(exitSignal, ps.Connects)

	// wait finished channel
	for i := 0; i < len(ps.Connects); i++ {
		<-finished
	}

	// Exit check signal
	exitSignal <- true

	// wait time (0.500 sec)
	time.Sleep(100 * time.Millisecond)

	// Exit Messages
	// Because it is Blocking.IO, you can not finish Input without input from the user.
	fmt.Fprintf(os.Stderr, "\n---\n%s\n", "Command exit. Please input Enter.")

	// Exit input
	exitInput <- true

	// Exit history
	for i := 0; i < len(ps.Connects); i++ {
		exitHistory <- true
	}

	// Add count
	ps.Count += 1
}

// pushSignal is send kill signal to session.
func (ps *pShell) pushKillSignal(exitSig chan bool, conns []*psConnect) (err error) {
	i := 0
	for {
		select {
		case <-ps.Signal:
			if i == 0 {
				for _, c := range conns {
					// send kill
					c.Kill()
				}
				i = 1
			}
		case <-exitSig:
			return
		}
	}
	return
}
<|endoftext|>


<commit_before>package stats

import (
	"math"
	"reflect"
	"sort"
	"testing"
)

func TestMin(t *testing.T) {
	m := Min([]float64{1.1, 2, 3, 4, 5})
	if m != 1.1 {
		t.Errorf("%.1f != %.1f", m, 1.1)
	}

	m = Min([]float64{10.534, 3, 5, 7, 9})
	if m != 3 {
		t.Errorf("%.1f != %.1f", m, 3)
	}

	m = Min([]float64{-5, 1, 5})
	if m != -5 {
		t.Errorf("%.1f != %.1f", m, -5)
	}
}

func TestMax(t *testing.T) {
	m := Max([]float64{1, 2, 3, 4, 5})
	if m != 5 {
		t.Errorf("%.1f != %.1f", m, 5)
	}

	m = Max([]float64{10.5, 3, 5, 7, 9})
	if m != 10.5 {
		t.Errorf("%.1f != %.1f", m, 10.5)
	}

	m = Max([]float64{-20, -1, -5.5})
	if m != -1 {
		t.Errorf("%.1f != %.1f", m, -1)
	}
}

func TestMean(t *testing.T) {
	m := Mean([]float64{1, 2, 3, 4, 5})
	if m != 3.0 {
		t.Errorf("%.1f != %.1f", m, 3.0)
	}

	m = Mean([]float64{1, 2, 3, 4, 5, 6})
	if m != 3.5 {
		t.Errorf("%.1f != %.1f", m, 3.5)
	}

	m = Mean([]float64{1})
	if m != 1.0 {
		t.Errorf("%.1f != %.1f", m, 1.0)
	}

	m = Mean([]float64{})
	if m != 0.0 {
		t.Errorf("%.1f != %.1f", m, 0.0)
	}
}

func TestMedian(t *testing.T) {
	m := Median([]float64{5, 3, 4, 2, 1})
	if m != 3.0 {
		t.Errorf("%.1f != %.1f", m, 3.0)
	}

	m = Median([]float64{6, 3, 2, 4, 5, 1})
	if m != 3.5 {
		t.Errorf("%.1f != %.1f", m, 3.5)
	}

	m = Median([]float64{1})
	if m != 1.0 {
		t.Errorf("%.1f != %.1f", m, 1.0)
	}

	m = Median([]float64{})
	if m != 0.0 {
		t.Errorf("%.1f != %.1f", m, 0.0)
	}
}

func TestMode(t *testing.T) {
	m := Mode([]float64{5, 3, 4, 2, 1})
	a := []float64{}
	if !reflect.DeepEqual(m, a) {
		t.Errorf("%.1f != %.1f", m, a)
	}

	m = Mode([]float64{5, 5, 3, 4, 2, 1})
	a = []float64{5}
	if !reflect.DeepEqual(m, a) {
		t.Errorf("%.1f != %.1f", m, a)
	}

	m = Mode([]float64{5, 5, 3, 3, 4, 2, 1})
	sort.Float64s(m)
	a = []float64{3, 5}
	if !reflect.DeepEqual(m, a) {
		t.Errorf("%.1f != %.1f", m, a)
	}

	m = Mode([]float64{5, 5, 3, 3, 4, 2, 1, 1, 1})
	a = []float64{1}
	if !reflect.DeepEqual(m, a) {
		t.Errorf("%.1f != %.1f", m, a)
	}
}

func TestSum(t *testing.T) {
	m := Sum([]float64{})
	if m != 0.0 {
		t.Errorf("%.1f != %.1f", m, 0.0)
	}

	m = Sum([]float64{1, 2, 3})
	if m != 6 {
		t.Errorf("%.1f != %.1f", m, 6)
	}

	m = Sum([]float64{1.0, 1.1, 1.2, 2.2})
	if m != 5.5 {
		t.Errorf("%.1f != %.1f", m, 5.5)
	}
}

func TestStdDev(t *testing.T) {
	m := Round(StdDev([]float64{1, 2, 3}), 2)
	if m != 0.82 {
		t.Errorf("%.10f != %.10f", m, 0.82)
	}

	m = Round(StdDev([]float64{-1, -2, -3.3}), 2)
	if m != 0.94 {
		t.Errorf("%.10f != %.10f", m, 0.94)
	}

	m = StdDev([]float64{})
	if m != 0.0 {
		t.Errorf("%.1f != %.1f", m, 0.0)
	}
}

func TestRound(t *testing.T) {
	m := Round(0.1111, 1)
	if m != 0.1 {
		t.Errorf("%.1f != %.1f", m, 0.1)
	}

	m = Round(-0.1111, 2)
	if m != -0.11 {
		t.Errorf("%.1f != %.1f", m, -0.11)
	}

	m = Round(5.3253, 3)
	if m != 5.325 {
		t.Errorf("%.1f != %.1f", m, 5.325)
	}

	m = Round(math.NaN(), 2)
	if !math.IsNaN(m) {
		t.Errorf("%.1f != %.1f", m, math.NaN())
	}
}
<commit_msg>Add another rounding test for whole numbers<commit_after>package stats

import (
	"math"
	"reflect"
	"sort"
	"testing"
)

func TestMin(t *testing.T) {
	m := Min([]float64{1.1, 2, 3, 4, 5})
	if m != 1.1 {
		t.Errorf("%.1f != %.1f", m, 1.1)
	}

	m = Min([]float64{10.534, 3, 5, 7, 9})
	if m != 3 {
		t.Errorf("%.1f != %.1f", m, 3)
	}

	m = Min([]float64{-5, 1, 5})
	if m != -5 {
		t.Errorf("%.1f != %.1f", m, -5)
	}
}

func TestMax(t *testing.T) {
	m := Max([]float64{1, 2, 3, 4, 5})
	if m != 5 {
		t.Errorf("%.1f != %.1f", m, 5)
	}

	m = Max([]float64{10.5, 3, 5, 7, 9})
	if m != 10.5 {
		t.Errorf("%.1f != %.1f", m, 10.5)
	}

	m = Max([]float64{-20, -1, -5.5})
	if m != -1 {
		t.Errorf("%.1f != %.1f", m, -1)
	}
}

func TestMean(t *testing.T) {
	m := Mean([]float64{1, 2, 3, 4, 5})
	if m != 3.0 {
		t.Errorf("%.1f != %.1f", m, 3.0)
	}

	m = Mean([]float64{1, 2, 3, 4, 5, 6})
	if m != 3.5 {
		t.Errorf("%.1f != %.1f", m, 3.5)
	}

	m = Mean([]float64{1})
	if m != 1.0 {
		t.Errorf("%.1f != %.1f", m, 1.0)
	}

	m = Mean([]float64{})
	if m != 0.0 {
		t.Errorf("%.1f != %.1f", m, 0.0)
	}
}

func TestMedian(t *testing.T) {
	m := Median([]float64{5, 3, 4, 2, 1})
	if m != 3.0 {
		t.Errorf("%.1f != %.1f", m, 3.0)
	}

	m = Median([]float64{6, 3, 2, 4, 5, 1})
	if m != 3.5 {
		t.Errorf("%.1f != %.1f", m, 3.5)
	}

	m = Median([]float64{1})
	if m != 1.0 {
		t.Errorf("%.1f != %.1f", m, 1.0)
	}

	m = Median([]float64{})
	if m != 0.0 {
		t.Errorf("%.1f != %.1f", m, 0.0)
	}
}

func TestMode(t *testing.T) {
	m := Mode([]float64{5, 3, 4, 2, 1})
	a := []float64{}
	if !reflect.DeepEqual(m, a) {
		t.Errorf("%.1f != %.1f", m, a)
	}

	m = Mode([]float64{5, 5, 3, 4, 2, 1})
	a = []float64{5}
	if !reflect.DeepEqual(m, a) {
		t.Errorf("%.1f != %.1f", m, a)
	}

	m = Mode([]float64{5, 5, 3, 3, 4, 2, 1})
	sort.Float64s(m)
	a = []float64{3, 5}
	if !reflect.DeepEqual(m, a) {
		t.Errorf("%.1f != %.1f", m, a)
	}

	m = Mode([]float64{5, 5, 3, 3, 4, 2, 1, 1, 1})
	a = []float64{1}
	if !reflect.DeepEqual(m, a) {
		t.Errorf("%.1f != %.1f", m, a)
	}
}

func TestSum(t *testing.T) {
	m := Sum([]float64{})
	if m != 0.0 {
		t.Errorf("%.1f != %.1f", m, 0.0)
	}

	m = Sum([]float64{1, 2, 3})
	if m != 6 {
		t.Errorf("%.1f != %.1f", m, 6)
	}

	m = Sum([]float64{1.0, 1.1, 1.2, 2.2})
	if m != 5.5 {
		t.Errorf("%.1f != %.1f", m, 5.5)
	}
}

func TestStdDev(t *testing.T) {
	m := Round(StdDev([]float64{1, 2, 3}), 2)
	if m != 0.82 {
		t.Errorf("%.10f != %.10f", m, 0.82)
	}

	m = Round(StdDev([]float64{-1, -2, -3.3}), 2)
	if m != 0.94 {
		t.Errorf("%.10f != %.10f", m, 0.94)
	}

	m = StdDev([]float64{})
	if m != 0.0 {
		t.Errorf("%.1f != %.1f", m, 0.0)
	}
}

func TestRound(t *testing.T) {
	m := Round(0.1111, 1)
	if m != 0.1 {
		t.Errorf("%.1f != %.1f", m, 0.1)
	}

	m = Round(-0.1111, 2)
	if m != -0.11 {
		t.Errorf("%.1f != %.1f", m, -0.11)
	}

	m = Round(5.3253, 3)
	if m != 5.325 {
		t.Errorf("%.1f != %.1f", m, 5.325)
	}

	m = Round(5.3253, 0)
	if m != 5 {
		t.Errorf("%.1f != %.1f", m, 5)
	}

	m = Round(math.NaN(), 2)
	if !math.IsNaN(m) {
		t.Errorf("%.1f != %.1f", m, math.NaN())
	}
}
<|endoftext|>


<commit_before>// Copyright ©2014 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package optimize

import (
	"math"

	"github.com/gonum/floats"
)

const (
	initialStepFactor = 1

	quadraticMinimumStepSize = 1e-3
	quadraticMaximumStepSize = 1.01
	quadraticThreshold       = 1e-12

	firstOrderMinimumStepSize = quadraticMinimumStepSize
	firstOrderMaximumStepSize = quadraticMaximumStepSize
)

// ConstantStepSize is a StepSizer that returns the same step size for
// every iteration.
type ConstantStepSize struct {
	Size float64
}

func (c ConstantStepSize) Init(l Location, dir []float64) float64 {
	return c.Size
}

func (c ConstantStepSize) StepSize(l Location, dir []float64) float64 {
	return c.Size
}

// QuadraticStepSize estimates the initial line search step size as the minimum
// of a quadratic that interpolates f(x_{k-1}), f(x_k) and ∇f_k⋅p_k.
// This is useful for line search methods that do not produce well-scaled
// descent directions, such as gradient descent or conjugate gradient methods.
// The step size is bounded away from zero.
type QuadraticStepSize struct {
	// If the relative change in the objective function is larger than
	// Threshold, the step size is estimated by quadratic interpolation,
	// otherwise it is set to 2*previous step size.
	// The default value is 1e-12.
	Threshold float64
	// The step size at the first iteration is estimated as InitialStepFactor / |g|_∞.
	// If InitialStepFactor is zero, it will be set to one.
	InitialStepFactor float64
	// The estimated step size is always greater than or equal to MinStepSize.
	// MinStepSize times GradientAbsTol should always be greater than machine epsilon.
	// If MinStepSize is zero, it will be set to 1e-3.
	MinStepSize float64
	// The estimated step size is always lower than or equal to MaxStepSize.
	// If MaxStepSize is zero, it will be set to 1.01.
	MaxStepSize float64

	fPrev        float64
	dirPrevNorm  float64
	projGradPrev float64
	xPrev        []float64
}

func (q *QuadraticStepSize) Init(l Location, dir []float64) (stepSize float64) {
	q.xPrev = resize(q.xPrev, len(l.X))
	if q.Threshold == 0 {
		q.Threshold = quadraticThreshold
	}
	if q.InitialStepFactor == 0 {
		q.InitialStepFactor = initialStepFactor
	}
	if q.MinStepSize == 0 {
		q.MinStepSize = quadraticMinimumStepSize
	}
	if q.MaxStepSize == 0 {
		q.MaxStepSize = quadraticMaximumStepSize
	}

	gNorm := floats.Norm(l.Gradient, math.Inf(1))
	stepSize = math.Max(q.MinStepSize, math.Min(q.InitialStepFactor/gNorm, q.MaxStepSize))

	q.fPrev = l.F
	q.dirPrevNorm = floats.Norm(dir, 2)
	q.projGradPrev = floats.Dot(l.Gradient, dir)
	copy(q.xPrev, l.X)
	return stepSize
}

func (q *QuadraticStepSize) StepSize(l Location, dir []float64) (stepSize float64) {
	stepSizePrev := floats.Distance(l.X, q.xPrev, 2) / q.dirPrevNorm
	projGrad := floats.Dot(l.Gradient, dir)

	stepSize = 2 * stepSizePrev
	if !floats.EqualWithinRel(q.fPrev, l.F, q.Threshold) {
		// Two consecutive function values are not relatively equal, so
		// computing the minimum of a quadratic interpolant might make sense

		df := (l.F - q.fPrev) / stepSizePrev
		quadTest := df - q.projGradPrev
		if quadTest > 0 {
			// There is a chance of approximating the function well by a
			// quadratic only if the finite difference (f_k-f_{k-1})/stepSizePrev
			// is larger than ∇f_{k-1}⋅p_{k-1}

			// Set the step size to the minimizer of the quadratic function that
			// interpolates f_{k-1}, ∇f_{k-1}⋅p_{k-1} and f_k
			stepSize = -q.projGradPrev * stepSizePrev / quadTest / 2
		}
	}
	// Bound the step size to lie in [MinStepSize, MaxStepSize]
	stepSize = math.Max(q.MinStepSize, math.Min(stepSize, q.MaxStepSize))

	q.fPrev = l.F
	q.dirPrevNorm = floats.Norm(dir, 2)
	q.projGradPrev = projGrad
	copy(q.xPrev, l.X)
	return stepSize
}

type FirstOrderStepSize struct {
	// The step size at the first iteration is estimated as InitialStepFactor / |g|_∞.
	// If InitialStepFactor is zero, it will be set to one.
	InitialStepFactor float64
	// The estimated step size is always greater than or equal to MinStepSize.
	// MinStepSize times GradientAbsTol should always be greater than machine epsilon.
	// If MinStepSize is zero, it will be set to 1e-3.
	MinStepSize float64
	// The estimated step size is always lower than or equal to MaxStepSize.
	// If MaxStepSize is zero, it will be set to 1.01.
	MaxStepSize float64

	dirPrevNorm  float64
	projGradPrev float64
	xPrev        []float64
}

func (fo *FirstOrderStepSize) Init(l Location, dir []float64) (stepSize float64) {
	fo.xPrev = resize(fo.xPrev, len(l.X))
	if fo.InitialStepFactor == 0 {
		fo.InitialStepFactor = initialStepFactor
	}
	if fo.MinStepSize == 0 {
		fo.MinStepSize = firstOrderMinimumStepSize
	}
	if fo.MaxStepSize == 0 {
		fo.MaxStepSize = firstOrderMaximumStepSize
	}

	gNorm := floats.Norm(l.Gradient, math.Inf(1))
	stepSize = math.Max(fo.MinStepSize, math.Min(fo.InitialStepFactor/gNorm, fo.MaxStepSize))

	fo.dirPrevNorm = floats.Norm(dir, 2)
	fo.projGradPrev = floats.Dot(l.Gradient, dir)
	copy(fo.xPrev, l.X)
	return stepSize
}

func (fo *FirstOrderStepSize) StepSize(l Location, dir []float64) (stepSize float64) {
	stepSizePrev := floats.Distance(l.X, fo.xPrev, 2) / fo.dirPrevNorm
	projGrad := floats.Dot(l.Gradient, dir)

	stepSize = stepSizePrev * fo.projGradPrev / projGrad
	stepSize = math.Max(fo.MinStepSize, math.Min(stepSize, fo.MaxStepSize))

	fo.dirPrevNorm = floats.Norm(dir, 2)
	fo.projGradPrev = floats.Dot(l.Gradient, dir)
	copy(fo.xPrev, l.X)
	return stepSize
}
<commit_msg>Add comment to FirstOrderStepSize<commit_after>// Copyright ©2014 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package optimize

import (
	"math"

	"github.com/gonum/floats"
)

const (
	initialStepFactor = 1

	quadraticMinimumStepSize = 1e-3
	quadraticMaximumStepSize = 1.01
	quadraticThreshold       = 1e-12

	firstOrderMinimumStepSize = quadraticMinimumStepSize
	firstOrderMaximumStepSize = quadraticMaximumStepSize
)

// ConstantStepSize is a StepSizer that returns the same step size for
// every iteration.
type ConstantStepSize struct {
	Size float64
}

func (c ConstantStepSize) Init(l Location, dir []float64) float64 {
	return c.Size
}

func (c ConstantStepSize) StepSize(l Location, dir []float64) float64 {
	return c.Size
}

// QuadraticStepSize estimates the initial line search step size as the minimum
// of a quadratic that interpolates f(x_{k-1}), f(x_k) and ∇f_k⋅p_k.
// This is useful for line search methods that do not produce well-scaled
// descent directions, such as gradient descent or conjugate gradient methods.
// The step size is bounded away from zero.
type QuadraticStepSize struct {
	// If the relative change in the objective function is larger than
	// Threshold, the step size is estimated by quadratic interpolation,
	// otherwise it is set to 2*previous step size.
	// The default value is 1e-12.
	Threshold float64
	// The step size at the first iteration is estimated as InitialStepFactor / |g|_∞.
	// If InitialStepFactor is zero, it will be set to one.
	InitialStepFactor float64
	// The estimated step size is always greater than or equal to MinStepSize.
	// MinStepSize times GradientAbsTol should always be greater than machine epsilon.
	// If MinStepSize is zero, it will be set to 1e-3.
	MinStepSize float64
	// The estimated step size is always lower than or equal to MaxStepSize.
	// If MaxStepSize is zero, it will be set to 1.01.
	MaxStepSize float64

	fPrev        float64
	dirPrevNorm  float64
	projGradPrev float64
	xPrev        []float64
}

func (q *QuadraticStepSize) Init(l Location, dir []float64) (stepSize float64) {
	q.xPrev = resize(q.xPrev, len(l.X))
	if q.Threshold == 0 {
		q.Threshold = quadraticThreshold
	}
	if q.InitialStepFactor == 0 {
		q.InitialStepFactor = initialStepFactor
	}
	if q.MinStepSize == 0 {
		q.MinStepSize = quadraticMinimumStepSize
	}
	if q.MaxStepSize == 0 {
		q.MaxStepSize = quadraticMaximumStepSize
	}

	gNorm := floats.Norm(l.Gradient, math.Inf(1))
	stepSize = math.Max(q.MinStepSize, math.Min(q.InitialStepFactor/gNorm, q.MaxStepSize))

	q.fPrev = l.F
	q.dirPrevNorm = floats.Norm(dir, 2)
	q.projGradPrev = floats.Dot(l.Gradient, dir)
	copy(q.xPrev, l.X)
	return stepSize
}

func (q *QuadraticStepSize) StepSize(l Location, dir []float64) (stepSize float64) {
	stepSizePrev := floats.Distance(l.X, q.xPrev, 2) / q.dirPrevNorm
	projGrad := floats.Dot(l.Gradient, dir)

	stepSize = 2 * stepSizePrev
	if !floats.EqualWithinRel(q.fPrev, l.F, q.Threshold) {
		// Two consecutive function values are not relatively equal, so
		// computing the minimum of a quadratic interpolant might make sense

		df := (l.F - q.fPrev) / stepSizePrev
		quadTest := df - q.projGradPrev
		if quadTest > 0 {
			// There is a chance of approximating the function well by a
			// quadratic only if the finite difference (f_k-f_{k-1})/stepSizePrev
			// is larger than ∇f_{k-1}⋅p_{k-1}

			// Set the step size to the minimizer of the quadratic function that
			// interpolates f_{k-1}, ∇f_{k-1}⋅p_{k-1} and f_k
			stepSize = -q.projGradPrev * stepSizePrev / quadTest / 2
		}
	}
	// Bound the step size to lie in [MinStepSize, MaxStepSize]
	stepSize = math.Max(q.MinStepSize, math.Min(stepSize, q.MaxStepSize))

	q.fPrev = l.F
	q.dirPrevNorm = floats.Norm(dir, 2)
	q.projGradPrev = projGrad
	copy(q.xPrev, l.X)
	return stepSize
}

// FirstOrderStepSize estimates the initial line search step size based on the
// assumption that the first-order change in the function will be the same as
// that obtained at the previous iteration. That is, the initial step size s^0_k
// is chosen so that
//   s^0_k ∇f_k⋅p_k = s_{k-1} ∇f_{k-1}⋅p_{k-1}
// This is useful for line search methods that do not produce well-scaled
// descent directions, such as gradient descent or conjugate gradient methods.
type FirstOrderStepSize struct {
	// The step size at the first iteration is estimated as InitialStepFactor / |g|_∞.
	// If InitialStepFactor is zero, it will be set to one.
	InitialStepFactor float64
	// The estimated step size is always greater than or equal to MinStepSize.
	// MinStepSize times GradientAbsTol should always be greater than machine epsilon.
	// If MinStepSize is zero, it will be set to 1e-3.
	MinStepSize float64
	// The estimated step size is always lower than or equal to MaxStepSize.
	// If MaxStepSize is zero, it will be set to 1.01.
	MaxStepSize float64

	dirPrevNorm  float64
	projGradPrev float64
	xPrev        []float64
}

func (fo *FirstOrderStepSize) Init(l Location, dir []float64) (stepSize float64) {
	fo.xPrev = resize(fo.xPrev, len(l.X))
	if fo.InitialStepFactor == 0 {
		fo.InitialStepFactor = initialStepFactor
	}
	if fo.MinStepSize == 0 {
		fo.MinStepSize = firstOrderMinimumStepSize
	}
	if fo.MaxStepSize == 0 {
		fo.MaxStepSize = firstOrderMaximumStepSize
	}

	gNorm := floats.Norm(l.Gradient, math.Inf(1))
	stepSize = math.Max(fo.MinStepSize, math.Min(fo.InitialStepFactor/gNorm, fo.MaxStepSize))

	fo.dirPrevNorm = floats.Norm(dir, 2)
	fo.projGradPrev = floats.Dot(l.Gradient, dir)
	copy(fo.xPrev, l.X)
	return stepSize
}

func (fo *FirstOrderStepSize) StepSize(l Location, dir []float64) (stepSize float64) {
	stepSizePrev := floats.Distance(l.X, fo.xPrev, 2) / fo.dirPrevNorm
	projGrad := floats.Dot(l.Gradient, dir)

	stepSize = stepSizePrev * fo.projGradPrev / projGrad
	stepSize = math.Max(fo.MinStepSize, math.Min(stepSize, fo.MaxStepSize))

	fo.dirPrevNorm = floats.Norm(dir, 2)
	fo.projGradPrev = floats.Dot(l.Gradient, dir)
	copy(fo.xPrev, l.X)
	return stepSize
}
<|endoftext|>


<commit_before>package storage

import (
	"crypto/sha512"
	"encoding/json"
	"fmt"
	"path"
	"time"

	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/s3"
)

type S3 struct {
	Bucket   *s3.Bucket
	hashFunc func(string) string
}

func NewS3(auth aws.Auth, region aws.Region, bucketName string) (*S3, error) {
	s := &S3{
		Bucket: s3.New(auth, region).Bucket(bucketName),
		hashFunc: func(s string) string {
			h := sha512.Sum384([]byte(s))
			return string(h[:])
		},
	}

	_, err := s.Bucket.GetBucketContents()
	if s3err, ok := err.(*s3.Error); ok && s3err.Code == "NoSuchBucket" {
		err = s.Bucket.PutBucket(s3.BucketOwnerFull)
	}

	return s, err
}

func (s *S3) saveKey(short, url string) (err error) {
	hashedShort := s.hashFunc(short)

	err = s.Bucket.Put(
		path.Join(hashedShort, "long"),
		[]byte(url),
		"text/plain",
		s3.BucketOwnerFull,
	)
	if err != nil {
		return err
	}

	changeLog, err := json.Marshal(
		struct {
			URL  string
			User string
		}{
			url,
			"TODO",
		},
	)
	if err != nil {
		return fmt.Errorf("unable to format change history: %v", err)
	}

	err = s.Bucket.Put(
		path.Join(hashedShort, "change_history", time.Now().Format(time.RFC3339Nano)),
		changeLog,
		"application/json",
		s3.BucketOwnerFull,
	)
	if err != nil {
		return fmt.Errorf("unable to save change history: %v", err)
	}

	return nil
}

func (s *S3) Save(url string) (string, error) {
	if _, err := validateURL(url); err != nil {
		return "", err
	}

	var (
		short string
		err   error
	)

	for i := 0; i < 10; i++ {
		short = getRandomString(8)

		if _, err = s.Bucket.GetKey(short); err != nil {
			return short, s.saveKey(short, url)
		}
	}

	return "", ErrShortNotSet
}

func (s *S3) SaveName(short string, url string) error {
	if err := validateShort(short); err != nil {
		return err
	}
	if _, err := validateURL(url); err != nil {
		return err
	}

	return s.saveKey(short, url)
}

func (s *S3) Load(short string) (string, error) {
	if err := validateShort(short); err != nil {
		return "", err
	}

	url, err := s.Bucket.Get(path.Join(s.hashFunc(short), "long"))
	if s3err, ok := err.(*s3.Error); ok && s3err.Code == "NoSuchKey" {
		return "", ErrShortNotSet
	}
	return string(url), err
}
<commit_msg>Whoops, base64 encode that Sha384<commit_after>package storage

import (
	"crypto/sha512"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"path"
	"time"

	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/s3"
)

type S3 struct {
	Bucket   *s3.Bucket
	hashFunc func(string) string
}

func NewS3(auth aws.Auth, region aws.Region, bucketName string) (*S3, error) {
	s := &S3{
		Bucket: s3.New(auth, region).Bucket(bucketName),
		hashFunc: func(s string) string {
			h := sha512.Sum384([]byte(s))
			return base64.StdEncoding.EncodeToString(h[:])
		},
	}

	_, err := s.Bucket.GetBucketContents()
	if s3err, ok := err.(*s3.Error); ok && s3err.Code == "NoSuchBucket" {
		err = s.Bucket.PutBucket(s3.BucketOwnerFull)
	}

	return s, err
}

func (s *S3) saveKey(short, url string) (err error) {
	hashedShort := s.hashFunc(short)

	err = s.Bucket.Put(
		path.Join(hashedShort, "long"),
		[]byte(url),
		"text/plain",
		s3.BucketOwnerFull,
	)
	if err != nil {
		return err
	}

	changeLog, err := json.Marshal(
		struct {
			URL  string
			User string
		}{
			url,
			"TODO",
		},
	)
	if err != nil {
		return fmt.Errorf("unable to format change history: %v", err)
	}

	err = s.Bucket.Put(
		path.Join(hashedShort, "change_history", time.Now().Format(time.RFC3339Nano)),
		changeLog,
		"application/json",
		s3.BucketOwnerFull,
	)
	if err != nil {
		return fmt.Errorf("unable to save change history: %v", err)
	}

	return nil
}

func (s *S3) Save(url string) (string, error) {
	if _, err := validateURL(url); err != nil {
		return "", err
	}

	var (
		short string
		err   error
	)

	for i := 0; i < 10; i++ {
		short = getRandomString(8)

		if _, err = s.Bucket.GetKey(short); err != nil {
			return short, s.saveKey(short, url)
		}
	}

	return "", ErrShortNotSet
}

func (s *S3) SaveName(short string, url string) error {
	if err := validateShort(short); err != nil {
		return err
	}
	if _, err := validateURL(url); err != nil {
		return err
	}

	return s.saveKey(short, url)
}

func (s *S3) Load(short string) (string, error) {
	if err := validateShort(short); err != nil {
		return "", err
	}

	url, err := s.Bucket.Get(path.Join(s.hashFunc(short), "long"))
	if s3err, ok := err.(*s3.Error); ok && s3err.Code == "NoSuchKey" {
		return "", ErrShortNotSet
	}
	return string(url), err
}
<|endoftext|>