<commit_before>package apps

import (
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/cozy/afero"
	"github.com/cozy/cozy-stack/pkg/magic"
	web_utils "github.com/cozy/cozy-stack/web/utils"
	"github.com/cozy/swift"
)

// FileServer interface defines a way to access and serve the application's
// data files.
type FileServer interface {
	Open(slug, version, shasum, file string) (io.ReadCloser, error)
	FilesList(slug, version, shasum string) ([]string, error)
	ServeFileContent(w http.ResponseWriter, req *http.Request,
		slug, version, shasum, file string) error
}

type swiftServer struct {
	c         *swift.Connection
	container string
}

type aferoServer struct {
	mkPath func(slug, version, shasum, file string) string
	fs     afero.Fs
}

type gzipReadCloser struct {
	gr *gzip.Reader
	cl io.Closer
}

// The Close method of gzip.Reader does not closes the underlying reader. This
// little wrapper does the closing.
func newGzipReadCloser(r io.ReadCloser) (io.ReadCloser, error) {
	gr, err := gzip.NewReader(r)
	if err != nil {
		return nil, err
	}
	return gzipReadCloser{gr: gr, cl: r}, nil
}

func (g gzipReadCloser) Read(b []byte) (int, error) {
	return g.gr.Read(b)
}

func (g gzipReadCloser) Close() error {
	err1 := g.gr.Close()
	err2 := g.cl.Close()
	if err1 != nil {
		return err1
	}
	if err2 != nil {
		return err2
	}
	return nil
}

// NewSwiftFileServer returns provides the apps.FileServer implementation
// using the swift backend as file server.
func NewSwiftFileServer(conn *swift.Connection, appsType AppType) FileServer {
	return &swiftServer{
		c:         conn,
		container: containerName(appsType),
	}
}

func (s *swiftServer) Open(slug, version, shasum, file string) (io.ReadCloser, error) {
	objName := s.makeObjectName(slug, version, shasum, file)
	f, h, err := s.c.ObjectOpen(s.container, objName, false, nil)
	if err != nil {
		return nil, wrapSwiftErr(err)
	}
	o := h.ObjectMetadata()
	if contentEncoding := o["content-encoding"]; contentEncoding == "gzip" {
		return newGzipReadCloser(f)
	}
	return f, nil
}

func (s *swiftServer) ServeFileContent(w http.ResponseWriter, req *http.Request, slug, version, shasum, file string) error {
	objName := s.makeObjectName(slug, version, shasum, file)
	f, h, err := s.c.ObjectOpen(s.container, objName, false, nil)
	if err != nil {
		return wrapSwiftErr(err)
	}
	defer f.Close()

	if checkETag := req.Header.Get("Cache-Control") == ""; checkETag {
		etag := fmt.Sprintf(`"%s"`, h["Etag"][:10])
		if web_utils.CheckPreconditions(w, req, etag) {
			return nil
		}
		w.Header().Set("Etag", etag)
	}

	var r io.Reader = f
	contentLength := h["Content-Length"]
	contentType := h["Content-Type"]
	o := h.ObjectMetadata()
	if contentEncoding := o["content-encoding"]; contentEncoding == "gzip" {
		if acceptGzipEncoding(req) {
			w.Header().Set("Content-Encoding", "gzip")
		} else {
			contentLength = o["original-content-length"]
			var gr *gzip.Reader
			gr, err = gzip.NewReader(f)
			if err != nil {
				return err
			}
			defer gr.Close()
			r = gr
		}
	}

	ext := path.Ext(file)
	if contentType == "" {
		contentType = magic.MIMETypeByExtension(ext)
	}
	if contentType == "text/html" {
		contentType = "text/html; charset=utf-8"
	} else if contentType == "text/xml" && ext == ".svg" {
		// override for files with text/xml content because of leading <?xml tag
		contentType = "image/svg+xml"
	}

	size, _ := strconv.ParseInt(contentLength, 10, 64)
	web_utils.ServeContent(w, req, contentType, size, r)
	return nil
}

func (s *swiftServer) makeObjectName(slug, version, shasum, file string) string {
	basepath := path.Join("/", slug, version)
	if shasum != "" {
		basepath += "-" + shasum
	}
	return path.Join(basepath, file)
}

func (s *swiftServer) FilesList(slug, version, shasum string) ([]string, error) {
	prefix := s.makeObjectName(slug, version, shasum, "") + "/"
	names, err := s.c.ObjectNamesAll(s.container, &swift.ObjectsOpts{
		Prefix: prefix,
	})
	if err != nil {
		return nil, err
	}
	filtered := names[:0]
	for _, n := range names {
		n = strings.TrimPrefix(n, prefix)
		if n != "" {
			filtered = append(filtered, n)
		}
	}
	return filtered, nil
}

// NewAferoFileServer returns a simple wrapper of the afero.Fs interface that
// provides the apps.FileServer interface.
//
// You can provide a makePath method to define how the file name should be
// created from the application's slug, version and file name. If not provided,
// the standard VFS concatenation (starting with vfs.WebappsDirName) is used.
func NewAferoFileServer(fs afero.Fs, makePath func(slug, version, shasum, file string) string) FileServer {
	if makePath == nil {
		makePath = defaultMakePath
	}
	return &aferoServer{
		mkPath: makePath,
		fs:     fs,
	}
}

func (s *aferoServer) Open(slug, version, shasum, file string) (io.ReadCloser, error) {
	isGzipped := true
	filepath := s.mkPath(slug, version, shasum, file)
	f, err := s.open(filepath + ".gz")
	if os.IsNotExist(err) {
		isGzipped = false
		f, err = s.open(filepath)
	}
	if err != nil {
		return nil, err
	}
	if isGzipped {
		return newGzipReadCloser(f)
	}
	return f, nil
}
func (s *aferoServer) open(filepath string) (io.ReadCloser, error) {
	return s.fs.Open(filepath)
}

func (s *aferoServer) ServeFileContent(w http.ResponseWriter, req *http.Request, slug, version, shasum, file string) error {
	filepath := s.mkPath(slug, version, shasum, file)
	return s.serveFileContent(w, req, filepath)
}
func (s *aferoServer) serveFileContent(w http.ResponseWriter, req *http.Request, filepath string) error {
	isGzipped := true
	rc, err := s.fs.Open(filepath + ".gz")
	if os.IsNotExist(err) {
		isGzipped = false
		rc, err = s.fs.Open(filepath)
	}
	if err != nil {
		return err
	}
	defer rc.Close()

	var content io.Reader
	var size int64
	if checkEtag := req.Header.Get("Cache-Control") == ""; checkEtag {
		var b []byte
		h := md5.New()
		r := io.TeeReader(rc, h)
		b, err = ioutil.ReadAll(r)
		if err != nil {
			return err
		}
		etag := fmt.Sprintf(`"%s"`, hex.EncodeToString(h.Sum(nil)))
		if web_utils.CheckPreconditions(w, req, etag) {
			return nil
		}
		w.Header().Set("Etag", etag)
		size = int64(len(b))
		content = bytes.NewReader(b)
	} else {
		size, err = rc.Seek(0, io.SeekEnd)
		if err != nil {
			return err
		}
		_, err = rc.Seek(0, io.SeekStart)
		if err != nil {
			return err
		}
		content = rc
	}

	if isGzipped {
		if acceptGzipEncoding(req) {
			w.Header().Set("Content-Encoding", "gzip")
		} else {
			var gr *gzip.Reader
			var b []byte
			gr, err = gzip.NewReader(content)
			if err != nil {
				return err
			}
			defer gr.Close()
			b, err = ioutil.ReadAll(gr)
			if err != nil {
				return err
			}
			size = int64(len(b))
			content = bytes.NewReader(b)
		}
	}

	contentType := magic.MIMETypeByExtension(path.Ext(filepath))
	if contentType == "text/html" {
		contentType = "text/html; charset=utf-8"
	}
	web_utils.ServeContent(w, req, contentType, size, content)
	return nil
}

func (s *aferoServer) FilesList(slug, version, shasum string) ([]string, error) {
	var names []string
	rootPath := s.mkPath(slug, version, shasum, "")
	err := afero.Walk(s.fs, rootPath, func(path string, infos os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !infos.IsDir() {
			name := strings.TrimPrefix(path, rootPath)
			name = strings.TrimSuffix(name, ".gz")
			names = append(names, name)
		}
		return nil
	})
	return names, err
}

func defaultMakePath(slug, version, shasum, file string) string {
	basepath := path.Join("/", slug, version)
	if shasum != "" {
		basepath += "-" + shasum
	}
	filepath := path.Join("/", file)
	return path.Join(basepath, filepath)
}

func acceptGzipEncoding(req *http.Request) bool {
	return strings.Contains(req.Header.Get("Accept-Encoding"), "gzip")
}

func containerName(appsType AppType) string {
	switch appsType {
	case Webapp:
		return "apps-web"
	case Konnector:
		return "apps-konnectors"
	}
	panic("Unknown AppType")
}

func wrapSwiftErr(err error) error {
	if err == swift.ObjectNotFound || err == swift.ContainerNotFound {
		return os.ErrNotExist
	}
	return err
}
<commit_msg>Fix the path of files in Swift for apps<commit_after>package apps

import (
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/cozy/afero"
	"github.com/cozy/cozy-stack/pkg/magic"
	web_utils "github.com/cozy/cozy-stack/web/utils"
	"github.com/cozy/swift"
)

// FileServer interface defines a way to access and serve the application's
// data files.
type FileServer interface {
	Open(slug, version, shasum, file string) (io.ReadCloser, error)
	FilesList(slug, version, shasum string) ([]string, error)
	ServeFileContent(w http.ResponseWriter, req *http.Request,
		slug, version, shasum, file string) error
}

type swiftServer struct {
	c         *swift.Connection
	container string
}

type aferoServer struct {
	mkPath func(slug, version, shasum, file string) string
	fs     afero.Fs
}

type gzipReadCloser struct {
	gr *gzip.Reader
	cl io.Closer
}

// The Close method of gzip.Reader does not closes the underlying reader. This
// little wrapper does the closing.
func newGzipReadCloser(r io.ReadCloser) (io.ReadCloser, error) {
	gr, err := gzip.NewReader(r)
	if err != nil {
		return nil, err
	}
	return gzipReadCloser{gr: gr, cl: r}, nil
}

func (g gzipReadCloser) Read(b []byte) (int, error) {
	return g.gr.Read(b)
}

func (g gzipReadCloser) Close() error {
	err1 := g.gr.Close()
	err2 := g.cl.Close()
	if err1 != nil {
		return err1
	}
	if err2 != nil {
		return err2
	}
	return nil
}

// NewSwiftFileServer returns provides the apps.FileServer implementation
// using the swift backend as file server.
func NewSwiftFileServer(conn *swift.Connection, appsType AppType) FileServer {
	return &swiftServer{
		c:         conn,
		container: containerName(appsType),
	}
}

func (s *swiftServer) Open(slug, version, shasum, file string) (io.ReadCloser, error) {
	objName := s.makeObjectName(slug, version, shasum, file)
	f, h, err := s.c.ObjectOpen(s.container, objName, false, nil)
	if err != nil {
		return nil, wrapSwiftErr(err)
	}
	o := h.ObjectMetadata()
	if contentEncoding := o["content-encoding"]; contentEncoding == "gzip" {
		return newGzipReadCloser(f)
	}
	return f, nil
}

func (s *swiftServer) ServeFileContent(w http.ResponseWriter, req *http.Request, slug, version, shasum, file string) error {
	objName := s.makeObjectName(slug, version, shasum, file)
	f, h, err := s.c.ObjectOpen(s.container, objName, false, nil)
	if err != nil {
		return wrapSwiftErr(err)
	}
	defer f.Close()

	if checkETag := req.Header.Get("Cache-Control") == ""; checkETag {
		etag := fmt.Sprintf(`"%s"`, h["Etag"][:10])
		if web_utils.CheckPreconditions(w, req, etag) {
			return nil
		}
		w.Header().Set("Etag", etag)
	}

	var r io.Reader = f
	contentLength := h["Content-Length"]
	contentType := h["Content-Type"]
	o := h.ObjectMetadata()
	if contentEncoding := o["content-encoding"]; contentEncoding == "gzip" {
		if acceptGzipEncoding(req) {
			w.Header().Set("Content-Encoding", "gzip")
		} else {
			contentLength = o["original-content-length"]
			var gr *gzip.Reader
			gr, err = gzip.NewReader(f)
			if err != nil {
				return err
			}
			defer gr.Close()
			r = gr
		}
	}

	ext := path.Ext(file)
	if contentType == "" {
		contentType = magic.MIMETypeByExtension(ext)
	}
	if contentType == "text/html" {
		contentType = "text/html; charset=utf-8"
	} else if contentType == "text/xml" && ext == ".svg" {
		// override for files with text/xml content because of leading <?xml tag
		contentType = "image/svg+xml"
	}

	size, _ := strconv.ParseInt(contentLength, 10, 64)
	web_utils.ServeContent(w, req, contentType, size, r)
	return nil
}

func (s *swiftServer) makeObjectName(slug, version, shasum, file string) string {
	basepath := path.Join(slug, version)
	if shasum != "" {
		basepath += "-" + shasum
	}
	return path.Join(basepath, file)
}

func (s *swiftServer) FilesList(slug, version, shasum string) ([]string, error) {
	prefix := s.makeObjectName(slug, version, shasum, "") + "/"
	names, err := s.c.ObjectNamesAll(s.container, &swift.ObjectsOpts{
		Prefix: prefix,
	})
	if err != nil {
		return nil, err
	}
	filtered := names[:0]
	for _, n := range names {
		n = strings.TrimPrefix(n, prefix)
		if n != "" {
			filtered = append(filtered, n)
		}
	}
	return filtered, nil
}

// NewAferoFileServer returns a simple wrapper of the afero.Fs interface that
// provides the apps.FileServer interface.
//
// You can provide a makePath method to define how the file name should be
// created from the application's slug, version and file name. If not provided,
// the standard VFS concatenation (starting with vfs.WebappsDirName) is used.
func NewAferoFileServer(fs afero.Fs, makePath func(slug, version, shasum, file string) string) FileServer {
	if makePath == nil {
		makePath = defaultMakePath
	}
	return &aferoServer{
		mkPath: makePath,
		fs:     fs,
	}
}

func (s *aferoServer) Open(slug, version, shasum, file string) (io.ReadCloser, error) {
	isGzipped := true
	filepath := s.mkPath(slug, version, shasum, file)
	f, err := s.open(filepath + ".gz")
	if os.IsNotExist(err) {
		isGzipped = false
		f, err = s.open(filepath)
	}
	if err != nil {
		return nil, err
	}
	if isGzipped {
		return newGzipReadCloser(f)
	}
	return f, nil
}
func (s *aferoServer) open(filepath string) (io.ReadCloser, error) {
	return s.fs.Open(filepath)
}

func (s *aferoServer) ServeFileContent(w http.ResponseWriter, req *http.Request, slug, version, shasum, file string) error {
	filepath := s.mkPath(slug, version, shasum, file)
	return s.serveFileContent(w, req, filepath)
}
func (s *aferoServer) serveFileContent(w http.ResponseWriter, req *http.Request, filepath string) error {
	isGzipped := true
	rc, err := s.fs.Open(filepath + ".gz")
	if os.IsNotExist(err) {
		isGzipped = false
		rc, err = s.fs.Open(filepath)
	}
	if err != nil {
		return err
	}
	defer rc.Close()

	var content io.Reader
	var size int64
	if checkEtag := req.Header.Get("Cache-Control") == ""; checkEtag {
		var b []byte
		h := md5.New()
		r := io.TeeReader(rc, h)
		b, err = ioutil.ReadAll(r)
		if err != nil {
			return err
		}
		etag := fmt.Sprintf(`"%s"`, hex.EncodeToString(h.Sum(nil)))
		if web_utils.CheckPreconditions(w, req, etag) {
			return nil
		}
		w.Header().Set("Etag", etag)
		size = int64(len(b))
		content = bytes.NewReader(b)
	} else {
		size, err = rc.Seek(0, io.SeekEnd)
		if err != nil {
			return err
		}
		_, err = rc.Seek(0, io.SeekStart)
		if err != nil {
			return err
		}
		content = rc
	}

	if isGzipped {
		if acceptGzipEncoding(req) {
			w.Header().Set("Content-Encoding", "gzip")
		} else {
			var gr *gzip.Reader
			var b []byte
			gr, err = gzip.NewReader(content)
			if err != nil {
				return err
			}
			defer gr.Close()
			b, err = ioutil.ReadAll(gr)
			if err != nil {
				return err
			}
			size = int64(len(b))
			content = bytes.NewReader(b)
		}
	}

	contentType := magic.MIMETypeByExtension(path.Ext(filepath))
	if contentType == "text/html" {
		contentType = "text/html; charset=utf-8"
	}
	web_utils.ServeContent(w, req, contentType, size, content)
	return nil
}

func (s *aferoServer) FilesList(slug, version, shasum string) ([]string, error) {
	var names []string
	rootPath := s.mkPath(slug, version, shasum, "")
	err := afero.Walk(s.fs, rootPath, func(path string, infos os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !infos.IsDir() {
			name := strings.TrimPrefix(path, rootPath)
			name = strings.TrimSuffix(name, ".gz")
			names = append(names, name)
		}
		return nil
	})
	return names, err
}

func defaultMakePath(slug, version, shasum, file string) string {
	basepath := path.Join("/", slug, version)
	if shasum != "" {
		basepath += "-" + shasum
	}
	filepath := path.Join("/", file)
	return path.Join(basepath, filepath)
}

func acceptGzipEncoding(req *http.Request) bool {
	return strings.Contains(req.Header.Get("Accept-Encoding"), "gzip")
}

func containerName(appsType AppType) string {
	switch appsType {
	case Webapp:
		return "apps-web"
	case Konnector:
		return "apps-konnectors"
	}
	panic("Unknown AppType")
}

func wrapSwiftErr(err error) error {
	if err == swift.ObjectNotFound || err == swift.ContainerNotFound {
		return os.ErrNotExist
	}
	return err
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"io"
	"os"

	"github.com/apprenda/kismatic/pkg/install"
	"github.com/apprenda/kismatic/pkg/util"
	"github.com/spf13/cobra"
)

// NewCmdUpgrade returns the upgrade command
func NewCmdUpgrade(out io.Writer) *cobra.Command {
	var planFile string
	cmd := &cobra.Command{
		Use:   "upgrade",
		Short: "upgrade your Kubernetes cluster",
		Run: func(cmd *cobra.Command, args []string) {
			cmd.Help()
		},
	}

	// Subcommands
	cmd.AddCommand(NewCmdUpgradeOffline(out, &planFile))
	addPlanFileFlag(cmd.PersistentFlags(), &planFile)
	return cmd
}

type upgradeOpts struct {
	generatedAssetsDir string
	verbose            bool
	outputFormat       string
	skipPreflight      bool
}

// NewCmdUpgradeOffline returns the command for running offline upgrades
func NewCmdUpgradeOffline(out io.Writer, planFile *string) *cobra.Command {
	opts := upgradeOpts{}
	cmd := cobra.Command{
		Use:   "offline",
		Short: "perform an offline upgrade of your Kubernetes cluster",
		RunE: func(cmd *cobra.Command, args []string) error {
			return doUpgradeOffline(out, *planFile, opts)
		},
	}
	cmd.Flags().StringVar(&opts.generatedAssetsDir, "generated-assets-dir", "generated", "path to the directory where assets generated during the installation process will be stored")
	cmd.Flags().BoolVar(&opts.verbose, "verbose", false, "enable verbose logging from the installation")
	cmd.Flags().StringVarP(&opts.outputFormat, "output", "o", "simple", "installation output format (options \"simple\"|\"raw\")")
	cmd.Flags().BoolVar(&opts.skipPreflight, "skip-preflight", false, "skip upgrade pre-flight checks")
	return &cmd
}

func doUpgradeOffline(out io.Writer, planFile string, opts upgradeOpts) error {
	planner := install.FilePlanner{File: planFile}
	executorOpts := install.ExecutorOptions{
		GeneratedAssetsDirectory: opts.generatedAssetsDir,
		RestartServices:          true,
		OutputFormat:             opts.outputFormat,
		Verbose:                  opts.verbose,
	}
	executor, err := install.NewExecutor(out, os.Stderr, executorOpts)
	if err != nil {
		return err
	}
	util.PrintHeader(out, "Computing upgrade plan", '=')

	// Read plan file
	if !planner.PlanExists() {
		util.PrettyPrintErr(out, "Reading plan file")
		return fmt.Errorf("plan file %q does not exist", planFile)
	}
	util.PrettyPrintOk(out, "Reading plan file")
	plan, err := planner.Read()
	if err != nil {
		util.PrettyPrintErr(out, "Reading plan file")
		return fmt.Errorf("error reading plan file %q: %v", planFile, err)
	}

	// Validate SSH connectivity to nodes
	if ok, errs := install.ValidatePlanSSHConnections(plan); !ok {
		util.PrettyPrintErr(out, "Validate SSH connectivity to nodes")
		util.PrintValidationErrors(out, errs)
		return fmt.Errorf("SSH connectivity validation errors found")
	}
	util.PrettyPrintOk(out, "Validate SSH connectivity to nodes")

	// Figure out which nodes to upgrade
	cv, err := install.ListVersions(plan)
	if err != nil {
		return fmt.Errorf("error listing cluster versions: %v", err)
	}
	var toUpgrade []install.ListableNode
	var toSkip []install.ListableNode
	for _, n := range cv.Nodes {
		if install.IsOlderVersion(n.Version) {
			toUpgrade = append(toUpgrade, n)
		} else {
			toSkip = append(toSkip, n)
		}
	}

	// Print the nodes that will be skipped
	if len(toSkip) > 0 {
		util.PrintHeader(out, "Skipping nodes", '=')
		for _, n := range toSkip {
			util.PrettyPrintOk(out, "- %q is at the target version %q", n.Node.IP, n.Version)
		}
		fmt.Fprintln(out)
	}

	// Print message if there's no work to do
	if len(toUpgrade) == 0 {
		fmt.Fprintln(out, "All nodes are at the target version. Skipping node upgrades.")
	} else {
		// Run upgrade preflight on the nodes that are to be UpgradeNodes
		if !opts.skipPreflight {
			util.PrintHeader(out, "Validating nodes", '=')
			if err := executor.RunUpgradePreFlightCheck(plan); err != nil {
				return fmt.Errorf("Upgrade preflight check failed: %v", err)
			}
		}
		// Run the upgrade on the nodes that need it
		if err := executor.UpgradeNodes(*plan, toUpgrade); err != nil {
			return fmt.Errorf("Failed to upgrade nodes: %v", err)
		}
		// validate on the master
		util.PrintHeader(out, "Validate Kubernetes Control Plane", '=')
		if err := executor.ValidateControlPlane(*plan); err != nil {
			return fmt.Errorf("Failed to validate kuberntes control plane: %v", err)
		}
	}

	if plan.ConfigureDockerRegistry() && plan.Cluster.DisconnectedInstallation {
		util.PrintHeader(out, "Upgrade Docker Registry", '=')
		if err := executor.UpgradeDockerRegistry(*plan); err != nil {
			return fmt.Errorf("Failed to upgrade docker registry: %v", err)
		}
	}

	// Upgrade the cluster services
	util.PrintHeader(out, "Upgrade Cluster Services", '=')
	if err := executor.UpgradeClusterServices(*plan); err != nil {
		return fmt.Errorf("Failed to upgrade cluster services: %v", err)
	}

	util.PrintHeader(out, "Smoke Test Cluster", '=')
	if err := executor.RunSmokeTest(plan); err != nil {
		return fmt.Errorf("Smoke test failed: %v", err)
	}

	fmt.Fprintln(out)
	util.PrintColor(out, util.Green, "Upgrade complete\n")
	fmt.Fprintln(out)
	return nil
}
<commit_msg>Add restart-services flag to upgrade<commit_after>package cli

import (
	"fmt"
	"io"
	"os"

	"github.com/apprenda/kismatic/pkg/install"
	"github.com/apprenda/kismatic/pkg/util"
	"github.com/spf13/cobra"
)

// NewCmdUpgrade returns the upgrade command
func NewCmdUpgrade(out io.Writer) *cobra.Command {
	var planFile string
	cmd := &cobra.Command{
		Use:   "upgrade",
		Short: "upgrade your Kubernetes cluster",
		Run: func(cmd *cobra.Command, args []string) {
			cmd.Help()
		},
	}

	// Subcommands
	cmd.AddCommand(NewCmdUpgradeOffline(out, &planFile))
	addPlanFileFlag(cmd.PersistentFlags(), &planFile)
	return cmd
}

type upgradeOpts struct {
	generatedAssetsDir string
	restartServices    bool
	verbose            bool
	outputFormat       string
	skipPreflight      bool
}

// NewCmdUpgradeOffline returns the command for running offline upgrades
func NewCmdUpgradeOffline(out io.Writer, planFile *string) *cobra.Command {
	opts := upgradeOpts{}
	cmd := cobra.Command{
		Use:   "offline",
		Short: "perform an offline upgrade of your Kubernetes cluster",
		RunE: func(cmd *cobra.Command, args []string) error {
			return doUpgradeOffline(out, *planFile, opts)
		},
	}
	cmd.Flags().StringVar(&opts.generatedAssetsDir, "generated-assets-dir", "generated", "path to the directory where assets generated during the installation process will be stored")
	cmd.Flags().BoolVar(&opts.restartServices, "restart-services", false, "force restart cluster services (Use with care)")
	cmd.Flags().BoolVar(&opts.verbose, "verbose", false, "enable verbose logging from the installation")
	cmd.Flags().StringVarP(&opts.outputFormat, "output", "o", "simple", "installation output format (options \"simple\"|\"raw\")")
	cmd.Flags().BoolVar(&opts.skipPreflight, "skip-preflight", false, "skip upgrade pre-flight checks")
	return &cmd
}

func doUpgradeOffline(out io.Writer, planFile string, opts upgradeOpts) error {
	planner := install.FilePlanner{File: planFile}
	executorOpts := install.ExecutorOptions{
		GeneratedAssetsDirectory: opts.generatedAssetsDir,
		RestartServices:          opts.restartServices,
		OutputFormat:             opts.outputFormat,
		Verbose:                  opts.verbose,
	}
	executor, err := install.NewExecutor(out, os.Stderr, executorOpts)
	if err != nil {
		return err
	}
	util.PrintHeader(out, "Computing upgrade plan", '=')

	// Read plan file
	if !planner.PlanExists() {
		util.PrettyPrintErr(out, "Reading plan file")
		return fmt.Errorf("plan file %q does not exist", planFile)
	}
	util.PrettyPrintOk(out, "Reading plan file")
	plan, err := planner.Read()
	if err != nil {
		util.PrettyPrintErr(out, "Reading plan file")
		return fmt.Errorf("error reading plan file %q: %v", planFile, err)
	}

	// Validate SSH connectivity to nodes
	if ok, errs := install.ValidatePlanSSHConnections(plan); !ok {
		util.PrettyPrintErr(out, "Validate SSH connectivity to nodes")
		util.PrintValidationErrors(out, errs)
		return fmt.Errorf("SSH connectivity validation errors found")
	}
	util.PrettyPrintOk(out, "Validate SSH connectivity to nodes")

	// Figure out which nodes to upgrade
	cv, err := install.ListVersions(plan)
	if err != nil {
		return fmt.Errorf("error listing cluster versions: %v", err)
	}
	var toUpgrade []install.ListableNode
	var toSkip []install.ListableNode
	for _, n := range cv.Nodes {
		if install.IsOlderVersion(n.Version) {
			toUpgrade = append(toUpgrade, n)
		} else {
			toSkip = append(toSkip, n)
		}
	}

	// Print the nodes that will be skipped
	if len(toSkip) > 0 {
		util.PrintHeader(out, "Skipping nodes", '=')
		for _, n := range toSkip {
			util.PrettyPrintOk(out, "- %q is at the target version %q", n.Node.IP, n.Version)
		}
		fmt.Fprintln(out)
	}

	// Print message if there's no work to do
	if len(toUpgrade) == 0 {
		fmt.Fprintln(out, "All nodes are at the target version. Skipping node upgrades.")
	} else {
		// Run upgrade preflight on the nodes that are to be UpgradeNodes
		if !opts.skipPreflight {
			util.PrintHeader(out, "Validating nodes", '=')
			if err := executor.RunUpgradePreFlightCheck(plan); err != nil {
				return fmt.Errorf("Upgrade preflight check failed: %v", err)
			}
		}
		// Run the upgrade on the nodes that need it
		if err := executor.UpgradeNodes(*plan, toUpgrade); err != nil {
			return fmt.Errorf("Failed to upgrade nodes: %v", err)
		}
		// validate on the master
		util.PrintHeader(out, "Validate Kubernetes Control Plane", '=')
		if err := executor.ValidateControlPlane(*plan); err != nil {
			return fmt.Errorf("Failed to validate kuberntes control plane: %v", err)
		}
	}

	if plan.ConfigureDockerRegistry() && plan.Cluster.DisconnectedInstallation {
		util.PrintHeader(out, "Upgrade Docker Registry", '=')
		if err := executor.UpgradeDockerRegistry(*plan); err != nil {
			return fmt.Errorf("Failed to upgrade docker registry: %v", err)
		}
	}

	// Upgrade the cluster services
	util.PrintHeader(out, "Upgrade Cluster Services", '=')
	if err := executor.UpgradeClusterServices(*plan); err != nil {
		return fmt.Errorf("Failed to upgrade cluster services: %v", err)
	}

	util.PrintHeader(out, "Smoke Test Cluster", '=')
	if err := executor.RunSmokeTest(plan); err != nil {
		return fmt.Errorf("Smoke test failed: %v", err)
	}

	fmt.Fprintln(out)
	util.PrintColor(out, util.Green, "Upgrade complete\n")
	fmt.Fprintln(out)
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package dump

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/crypto/ssh"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/klog/v2"
)

// logDumper gets all the nodes from a kubernetes cluster and dumps a well-known set of logs
type logDumper struct {
	sshClientFactory sshClientFactory

	artifactsDir string

	services     []string
	files        []string
	podSelectors []string
}

// NewLogDumper is the constructor for a logDumper
func NewLogDumper(sshConfig *ssh.ClientConfig, artifactsDir string) *logDumper {
	sshClientFactory := &sshClientFactoryImplementation{
		sshConfig: sshConfig,
	}

	d := &logDumper{
		sshClientFactory: sshClientFactory,
		artifactsDir:     artifactsDir,
	}

	d.services = []string{
		"node-problem-detector",
		"kubelet",
		"containerd",
		"docker",
		"kops-configuration",
		"protokube",
	}
	d.files = []string{
		"kube-apiserver",
		"kube-scheduler",
		"rescheduler",
		"kube-controller-manager",
		"etcd",
		"etcd-events",
		"etcd-cilium",
		"glbc",
		"cluster-autoscaler",
		"kube-addon-manager",
		"fluentd",
		"kube-proxy",
		"node-problem-detector",
		"cloud-init-output",
		"startupscript",
		"kern",
		"docker",
	}
	d.podSelectors = []string{
		"k8s-app=external-dns",
		"k8s-app=dns-controller",
	}

	return d
}

// DumpAllNodes connects to every node from kubectl get nodes and dumps the logs.
// additionalIPs holds IP addresses of instances found by the deployment tool;
// if the IPs are not found from kubectl get nodes, then these will be dumped also.
// This allows for dumping log on nodes even if they don't register as a kubernetes
// node, or if a node fails to register, or if the whole cluster fails to start.
func (d *logDumper) DumpAllNodes(ctx context.Context, nodes corev1.NodeList, additionalIPs []string) error {
	var dumped []*corev1.Node

	for i := range nodes.Items {
		if ctx.Err() != nil {
			log.Printf("stopping dumping nodes: %v", ctx.Err())
			return ctx.Err()
		}

		node := &nodes.Items[i]

		ip := ""
		for _, address := range node.Status.Addresses {
			if address.Type == "ExternalIP" {
				ip = address.Address
				break
			}
		}

		err := d.dumpNode(ctx, node.Name, ip)
		if err != nil {
			log.Printf("could not dump node %s (%s): %v", node.Name, ip, err)
		} else {
			dumped = append(dumped, node)
		}
	}

	notDumped := findInstancesNotDumped(additionalIPs, dumped)
	for _, ip := range notDumped {
		if ctx.Err() != nil {
			log.Printf("stopping dumping nodes: %v", ctx.Err())
			return ctx.Err()
		}

		log.Printf("dumping node not registered in kubernetes: %s", ip)
		err := d.dumpNode(ctx, ip, ip)
		if err != nil {
			log.Printf("error dumping node %s: %v", ip, err)
		}
	}

	return nil
}

// findInstancesNotDumped returns ips from the slice that do not appear as any address of the nodes
func findInstancesNotDumped(ips []string, dumped []*corev1.Node) []string {
	var notDumped []string
	dumpedAddresses := make(map[string]bool)
	for _, node := range dumped {
		for _, address := range node.Status.Addresses {
			dumpedAddresses[address.Address] = true
		}
	}

	for _, ip := range ips {
		if !dumpedAddresses[ip] {
			notDumped = append(notDumped, ip)
		}
	}
	return notDumped
}

// DumpNode connects to a node and dumps the logs.
func (d *logDumper) dumpNode(ctx context.Context, name string, ip string) error {
	if ip == "" {
		return fmt.Errorf("could not find address for %v, ", name)
	}

	log.Printf("Dumping node %s", name)

	n, err := d.connectToNode(ctx, name, ip)
	if err != nil {
		return fmt.Errorf("could not connect: %v", err)
	}

	// As long as we connect to the node we will not return an error;
	// a failure to collect a log (or even any logs at all) is not
	// considered an error in dumping the node.
	// TODO(justinsb): clean up / rationalize
	errors := n.dump(ctx)
	for _, e := range errors {
		log.Printf("error dumping node %s: %v", name, e)
	}

	if err := n.Close(); err != nil {
		log.Printf("error closing connection: %v", err)
	}

	return nil
}

// sshClient is an interface abstracting *ssh.Client, which allows us to test it
type sshClient interface {
	io.Closer

	// ExecPiped runs the command, piping stdout & stderr
	ExecPiped(ctx context.Context, command string, stdout io.Writer, stderr io.Writer) error
}

// sshClientFactory is an interface abstracting to a node over SSH
type sshClientFactory interface {
	Dial(ctx context.Context, host string) (sshClient, error)
}

// logDumperNode holds state for a particular node we are dumping
type logDumperNode struct {
	client sshClient
	dumper *logDumper

	dir string
}

// connectToNode makes an SSH connection to the node and returns a logDumperNode
func (d *logDumper) connectToNode(ctx context.Context, nodeName string, host string) (*logDumperNode, error) {
	client, err := d.sshClientFactory.Dial(ctx, host)
	if err != nil {
		return nil, fmt.Errorf("unable to SSH to %q: %v", host, err)
	}
	return &logDumperNode{
		client: client,
		dir:    filepath.Join(d.artifactsDir, nodeName),
		dumper: d,
	}, nil
}

// logDumperNode cleans up any state in the logDumperNode
func (n *logDumperNode) Close() error {
	return n.client.Close()
}

// dump captures the well-known set of logs
func (n *logDumperNode) dump(ctx context.Context) []error {
	if ctx.Err() != nil {
		return []error{ctx.Err()}
	}

	var errors []error

	// Capture kernel log
	if err := n.shellToFile(ctx, "sudo journalctl --output=short-precise -k", filepath.Join(n.dir, "kern.log")); err != nil {
		errors = append(errors, err)
	}

	// Capture full journal - needed so we can see e.g. disk mounts
	// This does duplicate the other files, but ensures we have all output
	if err := n.shellToFile(ctx, "sudo journalctl --output=short-precise", filepath.Join(n.dir, "journal.log")); err != nil {
		errors = append(errors, err)
	}

	// Capture logs from any systemd services in our list that are registered
	services, err := n.listSystemdUnits(ctx)
	if err != nil {
		errors = append(errors, fmt.Errorf("error listing systemd services: %v", err))
	}
	for _, s := range n.dumper.services {
		name := s + ".service"
		for _, service := range services {
			if service == name {
				if err := n.shellToFile(ctx, "sudo journalctl --output=cat -u "+name, filepath.Join(n.dir, s+".log")); err != nil {
					errors = append(errors, err)
				}
			}
		}
	}

	// Capture any file logs where the files exist
	fileList, err := n.findFiles(ctx, "/var/log")
	if err != nil {
		errors = append(errors, fmt.Errorf("error reading /var/log: %v", err))
	}
	for _, name := range n.dumper.files {
		prefix := "/var/log/" + name + ".log"
		for _, f := range fileList {
			if !strings.HasPrefix(f, prefix) {
				continue
			}
			if err := n.shellToFile(ctx, "sudo cat '"+strings.ReplaceAll(f, "'", "'\\''")+"'", filepath.Join(n.dir, filepath.Base(f))); err != nil {
				errors = append(errors, err)
			}
		}
	}

	for _, selector := range n.dumper.podSelectors {
		kv := strings.Split(selector, "=")
		app := kv[len(kv)-1]
		if err := n.shellToFile(ctx, "if command -v kubectl &> /dev/null; then kubectl logs -n kube-system --all-containers -l \""+selector+"\"; fi", filepath.Join(n.dir, app)); err != nil {
			errors = append(errors, err)
		}
	}

	return errors
}

// findFiles lists files under the specified directory (recursively)
func (n *logDumperNode) findFiles(ctx context.Context, dir string) ([]string, error) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := n.client.ExecPiped(ctx, "sudo find "+dir+" -print0", &stdout, &stderr)
	if err != nil {
		return nil, fmt.Errorf("error listing %q: %v", dir, err)
	}

	paths := []string{}
	for _, b := range bytes.Split(stdout.Bytes(), []byte{0}) {
		if len(b) == 0 {
			// Likely the last value
			continue
		}
		paths = append(paths, string(b))
	}
	return paths, nil
}

// listSystemdUnits returns the list of systemd units on the node
func (n *logDumperNode) listSystemdUnits(ctx context.Context) ([]string, error) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := n.client.ExecPiped(ctx, "sudo systemctl list-units -t service --no-pager --no-legend --all", &stdout, &stderr)
	if err != nil {
		return nil, fmt.Errorf("error listing systemd units: %v", err)
	}

	var services []string
	for _, line := range strings.Split(stdout.String(), "\n") {
		tokens := strings.Fields(line)
		if len(tokens) == 0 || tokens[0] == "" {
			continue
		}
		services = append(services, tokens[0])
	}
	return services, nil
}

// shellToFile executes a command and copies the output to a file
func (n *logDumperNode) shellToFile(ctx context.Context, command string, destPath string) error {
	if err := os.MkdirAll(filepath.Dir(destPath), 0755); err != nil {
		log.Printf("unable to mkdir on %q: %v", filepath.Dir(destPath), err)
	}

	f, err := os.Create(destPath)
	if err != nil {
		return fmt.Errorf("error creating file %q: %v", destPath, err)
	}
	defer f.Close()

	if err := n.client.ExecPiped(ctx, command, f, f); err != nil {
		return fmt.Errorf("error executing command %q: %v", command, err)
	}

	return nil
}

// sshClientImplementation is the default implementation of sshClient, binding to a *ssh.Client
type sshClientImplementation struct {
	client *ssh.Client
}

var _ sshClient = &sshClientImplementation{}

// ExecPiped implements sshClientImplementation::ExecPiped
func (s *sshClientImplementation) ExecPiped(ctx context.Context, cmd string, stdout io.Writer, stderr io.Writer) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	finished := make(chan error)
	go func() {
		session, err := s.client.NewSession()
		if err != nil {
			finished <- fmt.Errorf("error creating ssh session: %v", err)
			return
		}
		defer session.Close()

		klog.V(2).Infof("running SSH command: %v", cmd)

		session.Stdout = stdout
		session.Stderr = stderr

		finished <- session.Run(cmd)
	}()

	select {
	case <-ctx.Done():
		log.Print("closing SSH tcp connection due to context completion")

		// terminate the TCP connection to force a disconnect - we assume everyone is using the same context.
		// We could make this better by sending a signal on the session, waiting and then closing the session,
		// and only if we still haven't succeeded then closing the TCP connection.  This is sufficient for our
		// current usage though - and hopefully that logic will be implemented in the SSH package itself.
		s.Close()

		<-finished // Wait for cancellation
		return ctx.Err()

	case err := <-finished:
		return err
	}
}

// Close implements sshClientImplementation::Close
func (s *sshClientImplementation) Close() error {
	return s.client.Close()
}

// sshClientFactoryImplementation is the default implementation of sshClientFactory
type sshClientFactoryImplementation struct {
	sshConfig *ssh.ClientConfig
}

var _ sshClientFactory = &sshClientFactoryImplementation{}

// Dial implements sshClientFactory::Dial
func (f *sshClientFactoryImplementation) Dial(ctx context.Context, host string) (sshClient, error) {
	addr := host + ":22"
	d := net.Dialer{}
	conn, err := d.DialContext(ctx, "tcp", addr)
	if err != nil {
		return nil, err
	}

	// We have a TCP connection; we will force-close it to support context cancellation

	var client *ssh.Client
	finished := make(chan error)
	go func() {
		c, chans, reqs, err := ssh.NewClientConn(conn, addr, f.sshConfig)
		if err == nil {
			client = ssh.NewClient(c, chans, reqs)
		}
		finished <- err
	}()

	select {
	case <-ctx.Done():
		log.Print("cancelling SSH tcp connection due to context completion")
		conn.Close() // Close the TCP connection to force cancellation
		<-finished   // Wait for cancellation
		return nil, ctx.Err()
	case err := <-finished:
		if err != nil {
			return nil, err
		}
		return &sshClientImplementation{
			client: client,
		}, nil
	}
}
<commit_msg>Add .log suffix to dumped pod logs<commit_after>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package dump

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/crypto/ssh"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/klog/v2"
)

// logDumper gets all the nodes from a kubernetes cluster and dumps a well-known set of logs
type logDumper struct {
	sshClientFactory sshClientFactory

	artifactsDir string

	services     []string
	files        []string
	podSelectors []string
}

// NewLogDumper is the constructor for a logDumper
func NewLogDumper(sshConfig *ssh.ClientConfig, artifactsDir string) *logDumper {
	sshClientFactory := &sshClientFactoryImplementation{
		sshConfig: sshConfig,
	}

	d := &logDumper{
		sshClientFactory: sshClientFactory,
		artifactsDir:     artifactsDir,
	}

	d.services = []string{
		"node-problem-detector",
		"kubelet",
		"containerd",
		"docker",
		"kops-configuration",
		"protokube",
	}
	d.files = []string{
		"kube-apiserver",
		"kube-scheduler",
		"rescheduler",
		"kube-controller-manager",
		"etcd",
		"etcd-events",
		"etcd-cilium",
		"glbc",
		"cluster-autoscaler",
		"kube-addon-manager",
		"fluentd",
		"kube-proxy",
		"node-problem-detector",
		"cloud-init-output",
		"startupscript",
		"kern",
		"docker",
	}
	d.podSelectors = []string{
		"k8s-app=external-dns",
		"k8s-app=dns-controller",
	}

	return d
}

// DumpAllNodes connects to every node from kubectl get nodes and dumps the logs.
// additionalIPs holds IP addresses of instances found by the deployment tool;
// if the IPs are not found from kubectl get nodes, then these will be dumped also.
// This allows for dumping log on nodes even if they don't register as a kubernetes
// node, or if a node fails to register, or if the whole cluster fails to start.
func (d *logDumper) DumpAllNodes(ctx context.Context, nodes corev1.NodeList, additionalIPs []string) error {
	var dumped []*corev1.Node

	for i := range nodes.Items {
		if ctx.Err() != nil {
			log.Printf("stopping dumping nodes: %v", ctx.Err())
			return ctx.Err()
		}

		node := &nodes.Items[i]

		ip := ""
		for _, address := range node.Status.Addresses {
			if address.Type == "ExternalIP" {
				ip = address.Address
				break
			}
		}

		err := d.dumpNode(ctx, node.Name, ip)
		if err != nil {
			log.Printf("could not dump node %s (%s): %v", node.Name, ip, err)
		} else {
			dumped = append(dumped, node)
		}
	}

	notDumped := findInstancesNotDumped(additionalIPs, dumped)
	for _, ip := range notDumped {
		if ctx.Err() != nil {
			log.Printf("stopping dumping nodes: %v", ctx.Err())
			return ctx.Err()
		}

		log.Printf("dumping node not registered in kubernetes: %s", ip)
		err := d.dumpNode(ctx, ip, ip)
		if err != nil {
			log.Printf("error dumping node %s: %v", ip, err)
		}
	}

	return nil
}

// findInstancesNotDumped returns ips from the slice that do not appear as any address of the nodes
func findInstancesNotDumped(ips []string, dumped []*corev1.Node) []string {
	var notDumped []string
	dumpedAddresses := make(map[string]bool)
	for _, node := range dumped {
		for _, address := range node.Status.Addresses {
			dumpedAddresses[address.Address] = true
		}
	}

	for _, ip := range ips {
		if !dumpedAddresses[ip] {
			notDumped = append(notDumped, ip)
		}
	}
	return notDumped
}

// DumpNode connects to a node and dumps the logs.
func (d *logDumper) dumpNode(ctx context.Context, name string, ip string) error {
	if ip == "" {
		return fmt.Errorf("could not find address for %v, ", name)
	}

	log.Printf("Dumping node %s", name)

	n, err := d.connectToNode(ctx, name, ip)
	if err != nil {
		return fmt.Errorf("could not connect: %v", err)
	}

	// As long as we connect to the node we will not return an error;
	// a failure to collect a log (or even any logs at all) is not
	// considered an error in dumping the node.
	// TODO(justinsb): clean up / rationalize
	errors := n.dump(ctx)
	for _, e := range errors {
		log.Printf("error dumping node %s: %v", name, e)
	}

	if err := n.Close(); err != nil {
		log.Printf("error closing connection: %v", err)
	}

	return nil
}

// sshClient is an interface abstracting *ssh.Client, which allows us to test it
type sshClient interface {
	io.Closer

	// ExecPiped runs the command, piping stdout & stderr
	ExecPiped(ctx context.Context, command string, stdout io.Writer, stderr io.Writer) error
}

// sshClientFactory is an interface abstracting to a node over SSH
type sshClientFactory interface {
	Dial(ctx context.Context, host string) (sshClient, error)
}

// logDumperNode holds state for a particular node we are dumping
type logDumperNode struct {
	client sshClient
	dumper *logDumper

	dir string
}

// connectToNode makes an SSH connection to the node and returns a logDumperNode
func (d *logDumper) connectToNode(ctx context.Context, nodeName string, host string) (*logDumperNode, error) {
	client, err := d.sshClientFactory.Dial(ctx, host)
	if err != nil {
		return nil, fmt.Errorf("unable to SSH to %q: %v", host, err)
	}
	return &logDumperNode{
		client: client,
		dir:    filepath.Join(d.artifactsDir, nodeName),
		dumper: d,
	}, nil
}

// logDumperNode cleans up any state in the logDumperNode
func (n *logDumperNode) Close() error {
	return n.client.Close()
}

// dump captures the well-known set of logs
func (n *logDumperNode) dump(ctx context.Context) []error {
	if ctx.Err() != nil {
		return []error{ctx.Err()}
	}

	var errors []error

	// Capture kernel log
	if err := n.shellToFile(ctx, "sudo journalctl --output=short-precise -k", filepath.Join(n.dir, "kern.log")); err != nil {
		errors = append(errors, err)
	}

	// Capture full journal - needed so we can see e.g. disk mounts
	// This does duplicate the other files, but ensures we have all output
	if err := n.shellToFile(ctx, "sudo journalctl --output=short-precise", filepath.Join(n.dir, "journal.log")); err != nil {
		errors = append(errors, err)
	}

	// Capture logs from any systemd services in our list that are registered
	services, err := n.listSystemdUnits(ctx)
	if err != nil {
		errors = append(errors, fmt.Errorf("error listing systemd services: %v", err))
	}
	for _, s := range n.dumper.services {
		name := s + ".service"
		for _, service := range services {
			if service == name {
				if err := n.shellToFile(ctx, "sudo journalctl --output=cat -u "+name, filepath.Join(n.dir, s+".log")); err != nil {
					errors = append(errors, err)
				}
			}
		}
	}

	// Capture any file logs where the files exist
	fileList, err := n.findFiles(ctx, "/var/log")
	if err != nil {
		errors = append(errors, fmt.Errorf("error reading /var/log: %v", err))
	}
	for _, name := range n.dumper.files {
		prefix := "/var/log/" + name + ".log"
		for _, f := range fileList {
			if !strings.HasPrefix(f, prefix) {
				continue
			}
			if err := n.shellToFile(ctx, "sudo cat '"+strings.ReplaceAll(f, "'", "'\\''")+"'", filepath.Join(n.dir, filepath.Base(f))); err != nil {
				errors = append(errors, err)
			}
		}
	}

	for _, selector := range n.dumper.podSelectors {
		kv := strings.Split(selector, "=")
		logFile := fmt.Sprintf("%v.log", kv[len(kv)-1])
		if err := n.shellToFile(ctx, "if command -v kubectl &> /dev/null; then kubectl logs -n kube-system --all-containers -l \""+selector+"\"; fi", filepath.Join(n.dir, logFile)); err != nil {
			errors = append(errors, err)
		}
	}

	return errors
}

// findFiles lists files under the specified directory (recursively)
func (n *logDumperNode) findFiles(ctx context.Context, dir string) ([]string, error) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := n.client.ExecPiped(ctx, "sudo find "+dir+" -print0", &stdout, &stderr)
	if err != nil {
		return nil, fmt.Errorf("error listing %q: %v", dir, err)
	}

	paths := []string{}
	for _, b := range bytes.Split(stdout.Bytes(), []byte{0}) {
		if len(b) == 0 {
			// Likely the last value
			continue
		}
		paths = append(paths, string(b))
	}
	return paths, nil
}

// listSystemdUnits returns the list of systemd units on the node
func (n *logDumperNode) listSystemdUnits(ctx context.Context) ([]string, error) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := n.client.ExecPiped(ctx, "sudo systemctl list-units -t service --no-pager --no-legend --all", &stdout, &stderr)
	if err != nil {
		return nil, fmt.Errorf("error listing systemd units: %v", err)
	}

	var services []string
	for _, line := range strings.Split(stdout.String(), "\n") {
		tokens := strings.Fields(line)
		if len(tokens) == 0 || tokens[0] == "" {
			continue
		}
		services = append(services, tokens[0])
	}
	return services, nil
}

// shellToFile executes a command and copies the output to a file
func (n *logDumperNode) shellToFile(ctx context.Context, command string, destPath string) error {
	if err := os.MkdirAll(filepath.Dir(destPath), 0755); err != nil {
		log.Printf("unable to mkdir on %q: %v", filepath.Dir(destPath), err)
	}

	f, err := os.Create(destPath)
	if err != nil {
		return fmt.Errorf("error creating file %q: %v", destPath, err)
	}
	defer f.Close()

	if err := n.client.ExecPiped(ctx, command, f, f); err != nil {
		return fmt.Errorf("error executing command %q: %v", command, err)
	}

	return nil
}

// sshClientImplementation is the default implementation of sshClient, binding to a *ssh.Client
type sshClientImplementation struct {
	client *ssh.Client
}

var _ sshClient = &sshClientImplementation{}

// ExecPiped implements sshClientImplementation::ExecPiped
func (s *sshClientImplementation) ExecPiped(ctx context.Context, cmd string, stdout io.Writer, stderr io.Writer) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	finished := make(chan error)
	go func() {
		session, err := s.client.NewSession()
		if err != nil {
			finished <- fmt.Errorf("error creating ssh session: %v", err)
			return
		}
		defer session.Close()

		klog.V(2).Infof("running SSH command: %v", cmd)

		session.Stdout = stdout
		session.Stderr = stderr

		finished <- session.Run(cmd)
	}()

	select {
	case <-ctx.Done():
		log.Print("closing SSH tcp connection due to context completion")

		// terminate the TCP connection to force a disconnect - we assume everyone is using the same context.
		// We could make this better by sending a signal on the session, waiting and then closing the session,
		// and only if we still haven't succeeded then closing the TCP connection.  This is sufficient for our
		// current usage though - and hopefully that logic will be implemented in the SSH package itself.
		s.Close()

		<-finished // Wait for cancellation
		return ctx.Err()

	case err := <-finished:
		return err
	}
}

// Close implements sshClientImplementation::Close
func (s *sshClientImplementation) Close() error {
	return s.client.Close()
}

// sshClientFactoryImplementation is the default implementation of sshClientFactory
type sshClientFactoryImplementation struct {
	sshConfig *ssh.ClientConfig
}

var _ sshClientFactory = &sshClientFactoryImplementation{}

// Dial implements sshClientFactory::Dial
func (f *sshClientFactoryImplementation) Dial(ctx context.Context, host string) (sshClient, error) {
	addr := host + ":22"
	d := net.Dialer{}
	conn, err := d.DialContext(ctx, "tcp", addr)
	if err != nil {
		return nil, err
	}

	// We have a TCP connection; we will force-close it to support context cancellation

	var client *ssh.Client
	finished := make(chan error)
	go func() {
		c, chans, reqs, err := ssh.NewClientConn(conn, addr, f.sshConfig)
		if err == nil {
			client = ssh.NewClient(c, chans, reqs)
		}
		finished <- err
	}()

	select {
	case <-ctx.Done():
		log.Print("cancelling SSH tcp connection due to context completion")
		conn.Close() // Close the TCP connection to force cancellation
		<-finished   // Wait for cancellation
		return nil, ctx.Err()
	case err := <-finished:
		if err != nil {
			return nil, err
		}
		return &sshClientImplementation{
			client: client,
		}, nil
	}
}
<|endoftext|>


<commit_before>package logger

import (
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"github.com/fatih/color"
	"golang.org/x/crypto/ssh/terminal"
)

const (
	defaultTerminalWidth = 120

	logProcessDefaultProcessMsg      = "[RUNNING]"
	logProcessSuccessStatus          = "[OK]"
	logProcessFailedStatus           = "[FAILED]"
	logProcessTimeFormat             = "%5.2f sec "
	logProcessInlineProcessMsgFormat = "%s ..."

	logStateRightPartsSeparator = " "
)

var (
	indent = 0
)

func LogProcessInline(msg string, processFunc func() error) error {
	return logProcessInlineBase(msg, processFunc, colorizeStep, colorizeSuccess)
}

func LogServiceProcessInline(msg string, processFunc func() error) error {
	return logProcessInlineBase(msg, processFunc, colorizeService, colorizeService)
}

func LogProcess(msg, processMsg string, processFunc func() error) error {
	return logProcessBase(msg, processMsg, processFunc, colorizeStep, colorizeSuccess)
}

func LogServiceProcess(msg, processMsg string, processFunc func() error) error {
	return logProcessBase(msg, processMsg, processFunc, colorizeService, colorizeService)
}

func LogState(msg, state string) {
	logStateBase(msg, state, "", colorizeStep, colorizeService)
}

func LogServiceState(msg, state string) {
	logStateBase(msg, state, "", colorizeService, colorizeService)
}

func LogStep(msg string) {
	LogStepF("%s\n", msg)
}

func LogStepF(format string, args ...interface{}) {
	colorizeLogBaseF(os.Stdout, colorizeStep, format, args...)
}

func LogService(msg string) {
	LogServiceF("%s\n", msg)
}

func LogServiceF(format string, args ...interface{}) {
	colorizeLogBaseF(os.Stdout, colorizeService, format, args...)
}

func LogInfo(msg string) {
	LogInfoF("%s\n", msg)
}

func LogInfoF(format string, args ...interface{}) {
	colorizeLogBaseF(os.Stdout, colorizeInfo, format, args...)
}

func LogWarning(msg string) {
	LogWarningF("%s\n", msg)
}

func LogWarningF(format string, args ...interface{}) {
	colorizeLogBaseF(os.Stderr, colorizeWarning, format, args...)
}

func colorizeLogBaseF(w io.Writer, colorizeFunc func(string) string, format string, args ...interface{}) {
	logF(w, colorizeFunc(fmt.Sprintf(format, args...)))
}

func log(w io.Writer, msg string) {
	logF(w, "%s\n", msg)
}

func logF(w io.Writer, format string, args ...interface{}) {
	msg := fmt.Sprintf(format, args...)
	logBase(w, fmt.Sprintf("%s%s", logIndent(), msg))
}

func logBase(w io.Writer, msg string) {
	fmt.Fprintf(w, msg)
}

func logIndent() string {
	return strings.Repeat("  ", indent)
}

func withLogIndent(f func() error) error {
	indentUp()
	err := f()
	indentDown()

	return err
}

func indentUp() {
	indent += 1
}

func indentDown() {
	if indent == 0 {
		return
	}

	indent -= 1
}

func logProcessInlineBase(processMsg string, processFunc func() error, colorizeProcessMsgFunc, colorizeSuccessFunc func(string) string) error {
	processMsg = fmt.Sprintf(logProcessInlineProcessMsgFormat, processMsg)
	colorizeLogBaseF(os.Stdout, colorizeProcessMsgFunc, "%s", processMsg)

	resultStatus := logProcessSuccessStatus
	resultColorize := colorizeSuccessFunc
	start := time.Now()

	err := withLogIndent(processFunc)
	if err != nil {
		resultStatus = logProcessFailedStatus
		resultColorize = colorizeFail
	}

	elapsedSeconds := fmt.Sprintf(logProcessTimeFormat, time.Since(start).Seconds())

	rightPart := prepareLogStateRightPart(processMsg, resultStatus, elapsedSeconds, resultColorize)
	logBase(os.Stdout, fmt.Sprintf("%s\n", rightPart))

	return err
}

func logProcessBase(msg, processMsg string, processFunc func() error, colorizeMsgFunc, colorizeSuccessFunc func(string) string) error {
	if processMsg == "" {
		processMsg = logProcessDefaultProcessMsg
	}

	logStateBase(msg, processMsg, "", colorizeMsgFunc, colorizeSuccessFunc)

	start := time.Now()
	resultStatus := logProcessSuccessStatus

	err := withLogIndent(processFunc)

	elapsedSeconds := fmt.Sprintf(logProcessTimeFormat, time.Since(start).Seconds())

	if err != nil {
		resultStatus = logProcessFailedStatus
		logStateBase(msg, resultStatus, elapsedSeconds, colorizeFail, colorizeFail)

		return err
	}

	logStateBase(msg, resultStatus, elapsedSeconds, colorizeMsgFunc, colorizeSuccessFunc)

	return nil
}

func logStateBase(msg string, state, time string, colorizeLeftPartFunc, colorizeRightPartFunc func(string) string) {
	leftPart := prepareLogStateLeftPart(msg, state, time, colorizeLeftPartFunc)
	rightPart := prepareLogStateRightPart(msg, state, time, colorizeRightPartFunc)
	log(os.Stdout, fmt.Sprintf("%s%s", leftPart, rightPart))
}

func prepareLogStateLeftPart(msg, state, time string, colorizeFunc func(string) string) string {
	var result string

	spaceLength := availableTerminalLineSpace(state, timeOrStub(time))
	if spaceLength > 0 {
		if spaceLength > len(msg) {
			result = msg
		} else {
			result = msg[0:spaceLength]
		}
	} else {
		return ""
	}

	return colorizeFunc(result)
}

func prepareLogStateRightPart(msg, state, time string, colorizeFunc func(string) string) string {
	var result string
	spaceLength := availableTerminalLineSpace(msg)

	rightPartLength := len(state + timeOrStub(time) + logStateRightPartsSeparator)
	if spaceLength-rightPartLength > 0 {
		result += strings.Repeat(" ", spaceLength-rightPartLength)
	}

	var rightPart []string
	rightPart = append(rightPart, colorizeFunc(state))
	rightPart = append(rightPart, colorizeFunc(time))

	result += strings.Join(rightPart, logStateRightPartsSeparator)

	return result
}

func timeOrStub(time string) string {
	if time == "" {
		return fmt.Sprintf(logProcessTimeFormat, 0.0)
	}

	return time
}

func availableTerminalLineSpace(parts ...string) int {
	logIndentLength := len(logIndent())
	msgsLength := len(strings.Join(parts, " "))

	return terminalWidth() - logIndentLength - msgsLength
}

func terminalWidth() int {
	if terminal.IsTerminal(int(os.Stdout.Fd())) {
		w, _, err := terminal.GetSize(int(os.Stdout.Fd()))
		if err != nil {
			panic(err)
		}

		return w
	}

	return defaultTerminalWidth
}

func colorizeFail(msg string) string {
	return colorizeWarning(msg)
}

func colorizeSuccess(msg string) string {
	return colorize(msg, color.FgGreen, color.Bold)
}

func colorizeStep(msg string) string {
	return colorize(msg, color.FgYellow, color.Bold)
}

func colorizeService(msg string) string {
	return colorize(msg, color.FgWhite, color.Bold)
}

func colorizeInfo(msg string) string {
	return colorize(msg, color.FgBlue)
}

func colorizeWarning(msg string) string {
	return colorize(msg, color.FgRed, color.Bold)
}

func colorize(msg string, attributes ...color.Attribute) string {
	return color.New(attributes...).Sprint(msg)
}
<commit_msg>[go dapp] support multiline logs<commit_after>package logger

import (
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"github.com/fatih/color"
	"golang.org/x/crypto/ssh/terminal"
)

const (
	defaultTerminalWidth = 120

	logProcessDefaultProcessMsg      = "[RUNNING]"
	logProcessSuccessStatus          = "[OK]"
	logProcessFailedStatus           = "[FAILED]"
	logProcessTimeFormat             = "%5.2f sec "
	logProcessInlineProcessMsgFormat = "%s ..."

	logStateRightPartsSeparator = " "
)

var (
	indent = 0
)

func LogProcessInline(msg string, processFunc func() error) error {
	return logProcessInlineBase(msg, processFunc, colorizeStep, colorizeSuccess)
}

func LogServiceProcessInline(msg string, processFunc func() error) error {
	return logProcessInlineBase(msg, processFunc, colorizeService, colorizeService)
}

func LogProcess(msg, processMsg string, processFunc func() error) error {
	return logProcessBase(msg, processMsg, processFunc, colorizeStep, colorizeSuccess)
}

func LogServiceProcess(msg, processMsg string, processFunc func() error) error {
	return logProcessBase(msg, processMsg, processFunc, colorizeService, colorizeService)
}

func LogState(msg, state string) {
	logStateBase(msg, state, "", colorizeStep, colorizeService)
}

func LogServiceState(msg, state string) {
	logStateBase(msg, state, "", colorizeService, colorizeService)
}

func LogStep(msg string) {
	LogStepF("%s\n", msg)
}

func LogStepF(format string, args ...interface{}) {
	colorizeLogBaseF(os.Stdout, colorizeStep, format, args...)
}

func LogService(msg string) {
	LogServiceF("%s\n", msg)
}

func LogServiceF(format string, args ...interface{}) {
	colorizeLogBaseF(os.Stdout, colorizeService, format, args...)
}

func LogInfo(msg string) {
	LogInfoF("%s\n", msg)
}

func LogInfoF(format string, args ...interface{}) {
	colorizeLogBaseF(os.Stdout, colorizeInfo, format, args...)
}

func LogWarning(msg string) {
	LogWarningF("%s\n", msg)
}

func LogWarningF(format string, args ...interface{}) {
	colorizeLogBaseF(os.Stderr, colorizeWarning, format, args...)
}

func colorizeLogBaseF(w io.Writer, colorizeFunc func(string) string, format string, args ...interface{}) {
	var colorizeLines []string
	lines := strings.Split(fmt.Sprintf(format, args...), "\n")
	for _, line := range lines {
		if line == "" {
			colorizeLines = append(colorizeLines, line)
		} else {
			colorizeLines = append(colorizeLines, colorizeFunc(line))
		}
	}

	logF(w, strings.Join(colorizeLines, "\n"))
}

func log(w io.Writer, msg string) {
	logF(w, "%s\n", msg)
}

func logF(w io.Writer, format string, args ...interface{}) {
	msg := fmt.Sprintf(format, args...)

	var linesWithIndent []string
	lines := strings.Split(msg, "\n")
	for _, line := range lines {
		if line == "" {
			linesWithIndent = append(linesWithIndent, line)
		} else {
			linesWithIndent = append(linesWithIndent, fmt.Sprintf("%s%s", logIndent(), line))
		}
	}

	logBase(w, strings.Join(linesWithIndent, "\n"))
}

func logBase(w io.Writer, msg string) {
	fmt.Fprintf(w, msg)
}

func logIndent() string {
	return strings.Repeat("  ", indent)
}

func withLogIndent(f func() error) error {
	indentUp()
	err := f()
	indentDown()

	return err
}

func indentUp() {
	indent += 1
}

func indentDown() {
	if indent == 0 {
		return
	}

	indent -= 1
}

func logProcessInlineBase(processMsg string, processFunc func() error, colorizeProcessMsgFunc, colorizeSuccessFunc func(string) string) error {
	processMsg = fmt.Sprintf(logProcessInlineProcessMsgFormat, processMsg)
	colorizeLogBaseF(os.Stdout, colorizeProcessMsgFunc, "%s", processMsg)

	resultStatus := logProcessSuccessStatus
	resultColorize := colorizeSuccessFunc
	start := time.Now()

	err := withLogIndent(processFunc)
	if err != nil {
		resultStatus = logProcessFailedStatus
		resultColorize = colorizeFail
	}

	elapsedSeconds := fmt.Sprintf(logProcessTimeFormat, time.Since(start).Seconds())

	rightPart := prepareLogStateRightPart(processMsg, resultStatus, elapsedSeconds, resultColorize)
	logBase(os.Stdout, fmt.Sprintf("%s\n", rightPart))

	return err
}

func logProcessBase(msg, processMsg string, processFunc func() error, colorizeMsgFunc, colorizeSuccessFunc func(string) string) error {
	if processMsg == "" {
		processMsg = logProcessDefaultProcessMsg
	}

	logStateBase(msg, processMsg, "", colorizeMsgFunc, colorizeSuccessFunc)

	start := time.Now()
	resultStatus := logProcessSuccessStatus

	err := withLogIndent(processFunc)

	elapsedSeconds := fmt.Sprintf(logProcessTimeFormat, time.Since(start).Seconds())

	if err != nil {
		resultStatus = logProcessFailedStatus
		logStateBase(msg, resultStatus, elapsedSeconds, colorizeFail, colorizeFail)

		return err
	}

	logStateBase(msg, resultStatus, elapsedSeconds, colorizeMsgFunc, colorizeSuccessFunc)

	return nil
}

func logStateBase(msg string, state, time string, colorizeLeftPartFunc, colorizeRightPartFunc func(string) string) {
	leftPart := prepareLogStateLeftPart(msg, state, time, colorizeLeftPartFunc)
	rightPart := prepareLogStateRightPart(msg, state, time, colorizeRightPartFunc)
	log(os.Stdout, fmt.Sprintf("%s%s", leftPart, rightPart))
}

func prepareLogStateLeftPart(msg, state, time string, colorizeFunc func(string) string) string {
	var result string

	spaceLength := availableTerminalLineSpace(state, timeOrStub(time))
	if spaceLength > 0 {
		if spaceLength > len(msg) {
			result = msg
		} else {
			result = msg[0:spaceLength]
		}
	} else {
		return ""
	}

	return colorizeFunc(result)
}

func prepareLogStateRightPart(msg, state, time string, colorizeFunc func(string) string) string {
	var result string
	spaceLength := availableTerminalLineSpace(msg)

	rightPartLength := len(state + timeOrStub(time) + logStateRightPartsSeparator)
	if spaceLength-rightPartLength > 0 {
		result += strings.Repeat(" ", spaceLength-rightPartLength)
	}

	var rightPart []string
	rightPart = append(rightPart, colorizeFunc(state))
	rightPart = append(rightPart, colorizeFunc(time))

	result += strings.Join(rightPart, logStateRightPartsSeparator)

	return result
}

func timeOrStub(time string) string {
	if time == "" {
		return fmt.Sprintf(logProcessTimeFormat, 0.0)
	}

	return time
}

func availableTerminalLineSpace(parts ...string) int {
	logIndentLength := len(logIndent())
	msgsLength := len(strings.Join(parts, " "))

	return terminalWidth() - logIndentLength - msgsLength
}

func terminalWidth() int {
	if terminal.IsTerminal(int(os.Stdout.Fd())) {
		w, _, err := terminal.GetSize(int(os.Stdout.Fd()))
		if err != nil {
			panic(err)
		}

		return w
	}

	return defaultTerminalWidth
}

func colorizeFail(msg string) string {
	return colorizeWarning(msg)
}

func colorizeSuccess(msg string) string {
	return colorize(msg, color.FgGreen, color.Bold)
}

func colorizeStep(msg string) string {
	return colorize(msg, color.FgYellow, color.Bold)
}

func colorizeService(msg string) string {
	return colorize(msg, color.FgWhite, color.Bold)
}

func colorizeInfo(msg string) string {
	return colorize(msg, color.FgBlue)
}

func colorizeWarning(msg string) string {
	return colorize(msg, color.FgRed, color.Bold)
}

func colorize(msg string, attributes ...color.Attribute) string {
	return color.New(attributes...).Sprint(msg)
}
<|endoftext|>


<commit_before>// Package test contains test utilities
package test

import (
	"context"
	"fmt"
	"log"
	"net"
	"net/http"

	serverv2 "github.com/envoyproxy/go-control-plane/pkg/server/v2"
	serverv3 "github.com/envoyproxy/go-control-plane/pkg/server/v3"
	testv2 "github.com/envoyproxy/go-control-plane/pkg/test/v2"
	testv3 "github.com/envoyproxy/go-control-plane/pkg/test/v3"
	"google.golang.org/grpc"

	gcplogger "github.com/envoyproxy/go-control-plane/pkg/log"
)

const (
	grpcMaxConcurrentStreams = 1000000
)

// HTTPGateway is a custom implementation of [gRPC gateway](https://github.com/grpc-ecosystem/grpc-gateway)
// specialized to Envoy xDS API.
type HTTPGateway struct {
	// Log is an optional log for errors in response write
	Log gcplogger.Logger

	GatewayV2 serverv2.HTTPGateway

	GatewayV3 serverv3.HTTPGateway
}

// RunAccessLogServer starts an accesslog server.
func RunAccessLogServer(ctx context.Context, alsv2 *testv2.AccessLogService, alsv3 *testv3.AccessLogService, alsPort uint) {
	grpcServer := grpc.NewServer()
	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", alsPort))
	if err != nil {
		log.Fatal(err)
	}

	testv2.RegisterAccessLogServer(grpcServer, alsv2)
	log.Printf("access log server listening on %d\n", alsPort)

	go func() {
		if err = grpcServer.Serve(lis); err != nil {
			log.Println(err)
		}
	}()
	<-ctx.Done()

	grpcServer.GracefulStop()
}

// RunManagementServer starts an xDS server at the given port.
func RunManagementServer(ctx context.Context, srv2 serverv2.Server, srv3 serverv3.Server, port uint) {
	// gRPC golang library sets a very small upper bound for the number gRPC/h2
	// streams over a single TCP connection. If a proxy multiplexes requests over
	// a single connection to the management server, then it might lead to
	// availability problems.
	var grpcOptions []grpc.ServerOption
	grpcOptions = append(grpcOptions, grpc.MaxConcurrentStreams(grpcMaxConcurrentStreams))
	grpcServer := grpc.NewServer(grpcOptions...)

	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		log.Fatal(err)
	}

	testv2.RegisterServer(grpcServer, srv2)
	testv3.RegisterServer(grpcServer, srv3)

	log.Printf("management server listening on %d\n", port)
	go func() {
		if err = grpcServer.Serve(lis); err != nil {
			log.Println(err)
		}
	}()
	<-ctx.Done()

	grpcServer.GracefulStop()
}

// RunManagementGateway starts an HTTP gateway to an xDS server.
func RunManagementGateway(ctx context.Context, srv2 serverv2.Server, srv3 serverv3.Server, port uint, lg gcplogger.Logger) {
	log.Printf("gateway listening HTTP/1.1 on %d\n", port)
	server := &http.Server{
		Addr: fmt.Sprintf(":%d", port),
		Handler: &HTTPGateway{
			GatewayV2: serverv2.HTTPGateway{Log: lg, Server: srv2},
			GatewayV3: serverv3.HTTPGateway{Log: lg, Server: srv3},
			Log:       lg,
		},
	}
	go func() {
		if err := server.ListenAndServe(); err != nil {
			log.Println(err)
		}
	}()
	<-ctx.Done()

	// Cleanup our gateway if we receive a shutdown
	server.Shutdown(ctx)
}

func (h *HTTPGateway) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
	bytes, code, err := h.GatewayV2.ServeHTTP(req)
	if code == http.StatusNotFound {
		bytes, code, err = h.GatewayV3.ServeHTTP(req)
	}

	if err != nil {
		http.Error(resp, err.Error(), code)
		return
	}

	if bytes == nil {
		resp.WriteHeader(http.StatusNotModified)
		return
	}

	if _, err = resp.Write(bytes); err != nil && h.Log != nil {
		h.Log.Errorf("gateway error: %v", err)
	}
}
<commit_msg>fix: add v3 accesslog test server (#368)<commit_after>// Package test contains test utilities
package test

import (
	"context"
	"fmt"
	"log"
	"net"
	"net/http"

	serverv2 "github.com/envoyproxy/go-control-plane/pkg/server/v2"
	serverv3 "github.com/envoyproxy/go-control-plane/pkg/server/v3"
	testv2 "github.com/envoyproxy/go-control-plane/pkg/test/v2"
	testv3 "github.com/envoyproxy/go-control-plane/pkg/test/v3"
	"google.golang.org/grpc"

	gcplogger "github.com/envoyproxy/go-control-plane/pkg/log"
)

const (
	grpcMaxConcurrentStreams = 1000000
)

// HTTPGateway is a custom implementation of [gRPC gateway](https://github.com/grpc-ecosystem/grpc-gateway)
// specialized to Envoy xDS API.
type HTTPGateway struct {
	// Log is an optional log for errors in response write
	Log gcplogger.Logger

	GatewayV2 serverv2.HTTPGateway

	GatewayV3 serverv3.HTTPGateway
}

// RunAccessLogServer starts an accesslog server.
func RunAccessLogServer(ctx context.Context, alsv2 *testv2.AccessLogService, alsv3 *testv3.AccessLogService, alsPort uint) {
	grpcServer := grpc.NewServer()
	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", alsPort))
	if err != nil {
		log.Fatal(err)
	}

	testv2.RegisterAccessLogServer(grpcServer, alsv2)
	testv3.RegisterAccessLogServer(grpcServer, alsv3)
	log.Printf("access log server listening on %d\n", alsPort)

	go func() {
		if err = grpcServer.Serve(lis); err != nil {
			log.Println(err)
		}
	}()
	<-ctx.Done()

	grpcServer.GracefulStop()
}

// RunManagementServer starts an xDS server at the given port.
func RunManagementServer(ctx context.Context, srv2 serverv2.Server, srv3 serverv3.Server, port uint) {
	// gRPC golang library sets a very small upper bound for the number gRPC/h2
	// streams over a single TCP connection. If a proxy multiplexes requests over
	// a single connection to the management server, then it might lead to
	// availability problems.
	var grpcOptions []grpc.ServerOption
	grpcOptions = append(grpcOptions, grpc.MaxConcurrentStreams(grpcMaxConcurrentStreams))
	grpcServer := grpc.NewServer(grpcOptions...)

	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		log.Fatal(err)
	}

	testv2.RegisterServer(grpcServer, srv2)
	testv3.RegisterServer(grpcServer, srv3)

	log.Printf("management server listening on %d\n", port)
	go func() {
		if err = grpcServer.Serve(lis); err != nil {
			log.Println(err)
		}
	}()
	<-ctx.Done()

	grpcServer.GracefulStop()
}

// RunManagementGateway starts an HTTP gateway to an xDS server.
func RunManagementGateway(ctx context.Context, srv2 serverv2.Server, srv3 serverv3.Server, port uint, lg gcplogger.Logger) {
	log.Printf("gateway listening HTTP/1.1 on %d\n", port)
	server := &http.Server{
		Addr: fmt.Sprintf(":%d", port),
		Handler: &HTTPGateway{
			GatewayV2: serverv2.HTTPGateway{Log: lg, Server: srv2},
			GatewayV3: serverv3.HTTPGateway{Log: lg, Server: srv3},
			Log:       lg,
		},
	}
	go func() {
		if err := server.ListenAndServe(); err != nil {
			log.Println(err)
		}
	}()
	<-ctx.Done()

	// Cleanup our gateway if we receive a shutdown
	server.Shutdown(ctx)
}

func (h *HTTPGateway) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
	bytes, code, err := h.GatewayV2.ServeHTTP(req)
	if code == http.StatusNotFound {
		bytes, code, err = h.GatewayV3.ServeHTTP(req)
	}

	if err != nil {
		http.Error(resp, err.Error(), code)
		return
	}

	if bytes == nil {
		resp.WriteHeader(http.StatusNotModified)
		return
	}

	if _, err = resp.Write(bytes); err != nil && h.Log != nil {
		h.Log.Errorf("gateway error: %v", err)
	}
}
<|endoftext|>


<commit_before>package plugin

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/pkg/errors"
	"gopkg.in/urfave/cli.v1"
)

// CommandPlugin is definition of mkr plugin
var CommandPlugin = cli.Command{
	Name:        "plugin",
	Usage:       "Manage mackerel plugin",
	Description: `[WIP] Manage mackerel plugin`,
	Subcommands: []cli.Command{
		{
			Name:        "install",
			Usage:       "install mackerel plugin",
			Description: `WIP`,
			Action:      doPluginInstall,
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "prefix",
					Usage: "plugin install location",
				},
			},
		},
	},
	Hidden: true,
}

// main function for mkr plugin install
func doPluginInstall(c *cli.Context) error {
	argInstallTarget := c.Args().First()
	if argInstallTarget == "" {
		return fmt.Errorf("Specify install target")
	}

	_, err := parseInstallTarget(argInstallTarget)
	if err != nil {
		return errors.Wrap(err, "Failed to install plugin while setup plugin directory")
	}

	prefix, err := setupPluginDir(c.String("prefix"))
	if err != nil {
		return errors.Wrap(err, "Failed to install plugin while parsing install target")
	}

	// Create a work directory for downloading and extracting an artifact
	workdir, err := ioutil.TempDir(prefix, "work-")
	if err != nil {
		return errors.Wrap(err, "failed to install plugin while creating a work directory")
	}
	defer os.RemoveAll(workdir)

	fmt.Println("do plugin install [wip]")
	return nil
}

// Parse install target string passed from args
// example is below
// - mackerelio/mackerel-plugin-sample
// - mackerel-plugin-sample
// - mackerelio/mackerel-plugin-sample@v0.0.1
func parseInstallTarget(target string) (*installTarget, error) {
	it := &installTarget{}

	ownerRepoAndReleaseTag := strings.Split(target, "@")
	var ownerRepo string
	switch len(ownerRepoAndReleaseTag) {
	case 1:
		ownerRepo = ownerRepoAndReleaseTag[0]
	case 2:
		ownerRepo = ownerRepoAndReleaseTag[0]
		it.releaseTag = ownerRepoAndReleaseTag[1]
	default:
		return nil, fmt.Errorf("Install target is invalid: %s", target)
	}

	ownerAndRepo := strings.Split(ownerRepo, "/")
	switch len(ownerAndRepo) {
	case 1:
		it.pluginName = ownerAndRepo[0]
	case 2:
		it.owner = ownerAndRepo[0]
		it.repo = ownerAndRepo[1]
	default:
		return nil, fmt.Errorf("Install target is invalid: %s", target)
	}

	return it, nil
}

// Create a directory for plugin install
func setupPluginDir(prefix string) (string, error) {
	if prefix == "" {
		prefix = "/opt/mackerel-agent/plugins"
	}
	err := os.MkdirAll(filepath.Join(prefix, "bin"), 0755)
	if err != nil {
		return "", err
	}
	return prefix, nil
}

// Download plugin artifact from `url` to `workdir`,
// and returns downloaded filepath
func downloadPluginArtifact(url, workdir string) (fpath string, err error) {
	// Create request to download
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return "", err
	}
	req.Header.Set("User-Agent", "mkr-plugin-installer/0.0.0")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		err = fmt.Errorf("http response not OK. code: %d, url: %s", resp.StatusCode, url)
		return "", err
	}

	// fpath is filepath where artifact will be saved
	fpath = filepath.Join(workdir, path.Base(url))

	// download artifact
	file, err := os.OpenFile(fpath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return "", err
	}
	defer file.Close()

	_, err = io.Copy(file, resp.Body)
	if err != nil {
		return "", err
	}

	return fpath, nil
}

func looksLikePlugin(name string) bool {
	return strings.HasPrefix(name, "check-") || strings.HasPrefix(name, "mackerel-plugin-")
}

type installTarget struct {
	owner      string
	repo       string
	pluginName string
	releaseTag string
}

// Make artifact's download URL
func (it *installTarget) makeDownloadURL() (string, error) {
	if it.owner != "" && it.repo != "" {
		if it.releaseTag == "" {
			// TODO: Make latest release download URL by github API
			return "", fmt.Errorf("not implemented")
		}
		filename := fmt.Sprintf("%s_%s_%s.zip", it.repo, runtime.GOOS, runtime.GOARCH)
		return fmt.Sprintf("https://github.com/%s/%s/releases/download/%s/%s",
			it.owner, it.repo, it.releaseTag, filename), nil
	}
	// TODO: Make download URL by plugin registry
	return "", fmt.Errorf("not implemented")
}
<commit_msg>A downloaded artifact file does not need execution permission<commit_after>package plugin

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/pkg/errors"
	"gopkg.in/urfave/cli.v1"
)

// CommandPlugin is definition of mkr plugin
var CommandPlugin = cli.Command{
	Name:        "plugin",
	Usage:       "Manage mackerel plugin",
	Description: `[WIP] Manage mackerel plugin`,
	Subcommands: []cli.Command{
		{
			Name:        "install",
			Usage:       "install mackerel plugin",
			Description: `WIP`,
			Action:      doPluginInstall,
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "prefix",
					Usage: "plugin install location",
				},
			},
		},
	},
	Hidden: true,
}

// main function for mkr plugin install
func doPluginInstall(c *cli.Context) error {
	argInstallTarget := c.Args().First()
	if argInstallTarget == "" {
		return fmt.Errorf("Specify install target")
	}

	_, err := parseInstallTarget(argInstallTarget)
	if err != nil {
		return errors.Wrap(err, "Failed to install plugin while setup plugin directory")
	}

	prefix, err := setupPluginDir(c.String("prefix"))
	if err != nil {
		return errors.Wrap(err, "Failed to install plugin while parsing install target")
	}

	// Create a work directory for downloading and extracting an artifact
	workdir, err := ioutil.TempDir(prefix, "work-")
	if err != nil {
		return errors.Wrap(err, "failed to install plugin while creating a work directory")
	}
	defer os.RemoveAll(workdir)

	fmt.Println("do plugin install [wip]")
	return nil
}

// Parse install target string passed from args
// example is below
// - mackerelio/mackerel-plugin-sample
// - mackerel-plugin-sample
// - mackerelio/mackerel-plugin-sample@v0.0.1
func parseInstallTarget(target string) (*installTarget, error) {
	it := &installTarget{}

	ownerRepoAndReleaseTag := strings.Split(target, "@")
	var ownerRepo string
	switch len(ownerRepoAndReleaseTag) {
	case 1:
		ownerRepo = ownerRepoAndReleaseTag[0]
	case 2:
		ownerRepo = ownerRepoAndReleaseTag[0]
		it.releaseTag = ownerRepoAndReleaseTag[1]
	default:
		return nil, fmt.Errorf("Install target is invalid: %s", target)
	}

	ownerAndRepo := strings.Split(ownerRepo, "/")
	switch len(ownerAndRepo) {
	case 1:
		it.pluginName = ownerAndRepo[0]
	case 2:
		it.owner = ownerAndRepo[0]
		it.repo = ownerAndRepo[1]
	default:
		return nil, fmt.Errorf("Install target is invalid: %s", target)
	}

	return it, nil
}

// Create a directory for plugin install
func setupPluginDir(prefix string) (string, error) {
	if prefix == "" {
		prefix = "/opt/mackerel-agent/plugins"
	}
	err := os.MkdirAll(filepath.Join(prefix, "bin"), 0755)
	if err != nil {
		return "", err
	}
	return prefix, nil
}

// Download plugin artifact from `url` to `workdir`,
// and returns downloaded filepath
func downloadPluginArtifact(url, workdir string) (fpath string, err error) {
	// Create request to download
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return "", err
	}
	req.Header.Set("User-Agent", "mkr-plugin-installer/0.0.0")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		err = fmt.Errorf("http response not OK. code: %d, url: %s", resp.StatusCode, url)
		return "", err
	}

	// fpath is filepath where artifact will be saved
	fpath = filepath.Join(workdir, path.Base(url))

	// download artifact
	file, err := os.OpenFile(fpath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return "", err
	}
	defer file.Close()

	_, err = io.Copy(file, resp.Body)
	if err != nil {
		return "", err
	}

	return fpath, nil
}

func looksLikePlugin(name string) bool {
	return strings.HasPrefix(name, "check-") || strings.HasPrefix(name, "mackerel-plugin-")
}

type installTarget struct {
	owner      string
	repo       string
	pluginName string
	releaseTag string
}

// Make artifact's download URL
func (it *installTarget) makeDownloadURL() (string, error) {
	if it.owner != "" && it.repo != "" {
		if it.releaseTag == "" {
			// TODO: Make latest release download URL by github API
			return "", fmt.Errorf("not implemented")
		}
		filename := fmt.Sprintf("%s_%s_%s.zip", it.repo, runtime.GOOS, runtime.GOARCH)
		return fmt.Sprintf("https://github.com/%s/%s/releases/download/%s/%s",
			it.owner, it.repo, it.releaseTag, filename), nil
	}
	// TODO: Make download URL by plugin registry
	return "", fmt.Errorf("not implemented")
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"strings"
	"testing"

	"veyron.io/veyron/veyron2"
	"veyron.io/veyron/veyron2/ipc"
	"veyron.io/veyron/veyron2/naming"
	"veyron.io/veyron/veyron2/rt"
	"veyron.io/veyron/veyron2/services/mgmt/binary"
	"veyron.io/veyron/veyron2/services/mgmt/build"
	"veyron.io/veyron/veyron2/verror"
	"veyron.io/veyron/veyron2/vlog"

	"veyron.io/veyron/veyron/profiles"
)

var errInternalError = verror.Internalf("internal error")

type mock struct{}

func (mock) Build(ctx build.BuilderBuildContext, arch build.Architecture, opsys build.OperatingSystem) ([]byte, error) {
	vlog.VI(2).Infof("Build(%v, %v) was called", arch, opsys)
	iterator := ctx.RecvStream()
	for iterator.Advance() {
	}
	if err := iterator.Err(); err != nil {
		vlog.Errorf("Advance() failed: %v", err)
		return nil, errInternalError
	}
	return nil, nil
}

func (mock) Describe(_ ipc.ServerContext, name string) (binary.Description, error) {
	vlog.VI(2).Infof("Describe(%v) was called", name)
	return binary.Description{}, nil
}

type dispatcher struct{}

func startServer(runtime veyron2.Runtime, t *testing.T) (ipc.Server, naming.Endpoint) {
	server, err := runtime.NewServer()
	if err != nil {
		t.Fatalf("NewServer failed: %v", err)
	}
	endpoint, err := server.Listen(profiles.LocalListenSpec)
	if err != nil {
		t.Fatalf("Listen(%s) failed: %v", profiles.LocalListenSpec, err)
	}
	unpublished := ""
	if err := server.Serve(unpublished, build.BuilderServer(&mock{}), nil); err != nil {
		t.Fatalf("Serve(%v) failed: %v", unpublished, err)
	}
	return server, endpoint
}

func stopServer(t *testing.T, server ipc.Server) {
	if err := server.Stop(); err != nil {
		t.Errorf("Stop() failed: %v", err)
	}
}

func TestBuildClient(t *testing.T) {
	var err error
	runtime, err = rt.New()
	if err != nil {
		t.Fatalf("Unexpected error initializing runtime: %s", err)
	}
	server, endpoint := startServer(runtime, t)
	defer stopServer(t, server)

	cmd := root()
	var stdout, stderr bytes.Buffer
	cmd.Init(nil, &stdout, &stderr)

	// Test the 'Build' command.
	if err := cmd.Execute([]string{"build", naming.JoinAddressName(endpoint.String(), ""), "veyron.io/veyron/veyron/tools/build"}); err != nil {
		t.Fatalf("%v", err)
	}
	if expected, got := "", strings.TrimSpace(stdout.String()); got != expected {
		t.Errorf("Unexpected output from build: got %q, expected %q", got, expected)
	}
}
<commit_msg>veyron.io/veyron/veyron/tools/build:  Convert to verror2.<commit_after>package main

import (
	"bytes"
	"strings"
	"testing"

	"veyron.io/veyron/veyron2"
	"veyron.io/veyron/veyron2/ipc"
	"veyron.io/veyron/veyron2/naming"
	"veyron.io/veyron/veyron2/rt"
	"veyron.io/veyron/veyron2/services/mgmt/binary"
	"veyron.io/veyron/veyron2/services/mgmt/build"
	verror "veyron.io/veyron/veyron2/verror2"
	"veyron.io/veyron/veyron2/vlog"

	"veyron.io/veyron/veyron/profiles"
)

type mock struct{}

func (mock) Build(ctx build.BuilderBuildContext, arch build.Architecture, opsys build.OperatingSystem) ([]byte, error) {
	vlog.VI(2).Infof("Build(%v, %v) was called", arch, opsys)
	iterator := ctx.RecvStream()
	for iterator.Advance() {
	}
	if err := iterator.Err(); err != nil {
		vlog.Errorf("Advance() failed: %v", err)
		return nil, verror.Make(verror.Internal, ctx)
	}
	return nil, nil
}

func (mock) Describe(_ ipc.ServerContext, name string) (binary.Description, error) {
	vlog.VI(2).Infof("Describe(%v) was called", name)
	return binary.Description{}, nil
}

type dispatcher struct{}

func startServer(runtime veyron2.Runtime, t *testing.T) (ipc.Server, naming.Endpoint) {
	server, err := runtime.NewServer()
	if err != nil {
		t.Fatalf("NewServer failed: %v", err)
	}
	endpoint, err := server.Listen(profiles.LocalListenSpec)
	if err != nil {
		t.Fatalf("Listen(%s) failed: %v", profiles.LocalListenSpec, err)
	}
	unpublished := ""
	if err := server.Serve(unpublished, build.BuilderServer(&mock{}), nil); err != nil {
		t.Fatalf("Serve(%v) failed: %v", unpublished, err)
	}
	return server, endpoint
}

func stopServer(t *testing.T, server ipc.Server) {
	if err := server.Stop(); err != nil {
		t.Errorf("Stop() failed: %v", err)
	}
}

func TestBuildClient(t *testing.T) {
	var err error
	runtime, err = rt.New()
	if err != nil {
		t.Fatalf("Unexpected error initializing runtime: %s", err)
	}
	server, endpoint := startServer(runtime, t)
	defer stopServer(t, server)

	cmd := root()
	var stdout, stderr bytes.Buffer
	cmd.Init(nil, &stdout, &stderr)

	// Test the 'Build' command.
	if err := cmd.Execute([]string{"build", naming.JoinAddressName(endpoint.String(), ""), "veyron.io/veyron/veyron/tools/build"}); err != nil {
		t.Fatalf("%v", err)
	}
	if expected, got := "", strings.TrimSpace(stdout.String()); got != expected {
		t.Errorf("Unexpected output from build: got %q, expected %q", got, expected)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"time"
	"os"
	"strconv"
	"errors"
	"io/ioutil"
	"github.com/cs3238-tsuzu/popcon/file_manager"
	"encoding/json"
)

// Database manager for Contest and Onlinejudge

//import "errors"

const SubmissionDir = "./submissions/"

type SubmissionStatus int64

const (
    InQueue SubmissionStatus = 0
    Judging SubmissionStatus = 1
    Accepted SubmissionStatus = 2
    WrongAnswer SubmissionStatus = 3
    TimeLimitExceeded SubmissionStatus = 4
    MemoryLimitExceeded SubmissionStatus = 5
    RuntimeError SubmissionStatus = 7
    CompileError SubmissionStatus = 8
    InternalError SubmissionStatus = 9
)

type Submission struct {
    Sid int64 `db:"pk" default:""`
    Pid int64 `default:""` //index
    Cid int64 `default:""` //index, onlinejudge->0
    Iid int64 `default:""` //index
    Lang int64 `default:""`
    Time int64 `default:""`
    Mem int64 `default:""`
    Score int64 `default:""`
    SubmitTime int64 `default:""` //提出日時
    Status SubmissionStatus `default:""` //index
    Prog uint64 `default:""` //テストケースの進捗状況(完了数<<32 & 全体数)
}

func (dm *DatabaseManager) CreateSubmissionTable() error {
    err := dm.db.CreateTableIfNotExists(&Submission{})

    if err != nil {
        return err
    }

    dm.db.CreateIndex(&Submission{}, "pid")
    dm.db.CreateIndex(&Submission{}, "cid")
    dm.db.CreateIndex(&Submission{}, "iid")
    dm.db.CreateIndex(&Submission{}, "status")

    return nil
}

func (dm *DatabaseManager) SubmissionNew(pid, cid, iid, lang int64, code string) (int64, error) {
    sm := Submission {
        Pid: pid,
        Cid: cid,
        Iid: iid,
        Lang: lang,
        Time: 0,
        Mem: 0,
        Score: 0,
        SubmitTime: time.Now().Unix(),
        Status: InQueue,
        Prog: 0,
    }

    id, err := dm.db.Insert(&sm)

    if err != nil {
        return 0, err
    }

    err = os.MkdirAll(SubmissionDir + strconv.FormatInt(id, 10), 0644)

    if err != nil {
        //SubmissionDelete

        return 0, err
    }

    fp, err := os.OpenFile(SubmissionDir + strconv.FormatInt(id, 10) + "/msg", os.O_WRONLY | os.O_CREATE, 0644)

    if err != nil {
        return 0, err
    }

    fp.Close()

    fp, err = os.OpenFile(SubmissionDir + strconv.FormatInt(id, 10) + "/case", os.O_WRONLY | os.O_CREATE, 0644)

    if err != nil {
        return 0, err
    }

    fp.Close()

    fp, err = os.OpenFile(SubmissionDir + strconv.FormatInt(id, 10) + "/code", os.O_WRONLY | os.O_CREATE, 0644)

    if err != nil {
        return 0, err
    }

    defer fp.Close()

    _, err = fp.Write([]byte(code))

    if err != nil {
        return 0, err
    }

    return id, nil
}

func (dm *DatabaseManager) SubmissionRemove(sid int64) error {
    sm := Submission{Sid: sid}
    _, err := dm.db.Delete(&sm)

    if err != nil {
        return err
    }

    fm1, _ := FileManager.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/case", os.O_RDONLY, true)
    fm2, _ := FileManager.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/msg", os.O_RDONLY, true)

    defer func () {
        fm1.Close()
        fm2.Close()
    }()

    return os.RemoveAll(SubmissionDir + strconv.FormatInt(sid, 10))
}

func (dm *DatabaseManager) SubmissionFind(sid int64) (*Submission, error) {
    var results []Submission

    err := dm.db.Select(&results, dm.db.Where("sid", "=", sid))

    if err != nil {
        return nil, err
    }

    if len(results) == 0 {
        return nil, errors.New("Unknown submission")
    }

    return &results[0], nil
}

func (dm *DatabaseManager) SubmissionUpdate(sid, time, mem int64, status SubmissionStatus, fin, all int) error {
    sm, err := dm.SubmissionFind(sid)

    if err != nil {
        return err
    }

    sm.Time = time
    sm.Mem = mem
    sm.Status = status
    sm.Prog = uint64(fin) << 32 & uint64(all)

    _, err = dm.db.Update(&sm)

    return err
}

func (dm *DatabaseManager) SubmissionGetCode(sid int64) (*string, error) {
    fp, err := os.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/code", os.O_RDONLY, 0644)

    if err != nil {
        return nil, err
    }

    defer fp.Close()

    b, err := ioutil.ReadAll(fp)

    if err != nil {
        return nil, err
    }

    str := string(b)

    return &str, nil
}

func (dm *DatabaseManager) SubmissionGetMsg(sid int64) string {
    fm, err := FileManager.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/msg", os.O_RDONLY, false)

    if err != nil {
        return ""
    }

    defer fm.Close()

    b, err := ioutil.ReadAll(fm)

    if err != nil {
        return ""
    }

    return string(b)
}

func (dm *DatabaseManager) SubmissionSetMsg(sid int64, msg string) error {
    fm, err := FileManager.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/msg", os.O_WRONLY | os.O_CREATE, true)

    if err != nil {
        return err
    }

    defer fm.Close()

    _, err = fm.Write([]byte(msg))

    return err
}

func (dm *DatabaseManager) SubmissionGetCase(sid int64) (*map[int64]SubmissionStatus, error) {
    fm, err := FileManager.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/case", os.O_RDONLY, false)

    if err != nil {
        return nil, err
    }

    defer fm.Close()

    b, err := ioutil.ReadAll(fm)

    if err != nil {
        return nil, err
    }

    var ss map[int64]SubmissionStatus
    err = json.Unmarshal(b, &ss)

    if err != nil {
        return nil, err
    }

    return &ss, nil
}

func (dm *DatabaseManager) SubmissionSetCase(sid int64, ss map[int64]SubmissionStatus) error {
    fm, err := FileManager.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/case", os.O_RDONLY, false)

    if err != nil {
        return err
    }

    defer fm.Close()

    b, err := json.Marshal(ss)

    if err != nil {
        return err
    }

    _, err = fm.Write(b)

    return err
}<commit_msg>Mon Jul 18 22:19:10 JST 2016<commit_after>package main

import (
	"time"
	"os"
	"strconv"
	"errors"
	"io/ioutil"
	"github.com/cs3238-tsuzu/popcon/file_manager"
	"encoding/json"
)

// Database manager for Contest and Onlinejudge

//import "errors"

const SubmissionDir = "./submissions/"

type SubmissionStatus int64

const (
    InQueue SubmissionStatus = 0
    Judging SubmissionStatus = 1
    Accepted SubmissionStatus = 2
    WrongAnswer SubmissionStatus = 3
    TimeLimitExceeded SubmissionStatus = 4
    MemoryLimitExceeded SubmissionStatus = 5
    RuntimeError SubmissionStatus = 7
    CompileError SubmissionStatus = 8
    InternalError SubmissionStatus = 9
)

type Submission struct {
    Sid int64 `db:"pk" default:""`
    Pid int64 `default:""` //index
    Cid int64 `default:""` //index, onlinejudge->0
    Iid int64 `default:""` //index
    Lang int64 `default:""`
    Time int64 `default:""`
    Mem int64 `default:""`
    Score int64 `default:""`
    SubmitTime int64 `default:""` //提出日時
    Status int64 `default:""` //index
    Prog uint64 `default:""` //テストケースの進捗状況(完了数<<32 & 全体数)
}

func (dm *DatabaseManager) CreateSubmissionTable() error {
    err := dm.db.CreateTableIfNotExists(&Submission{})

    if err != nil {
        return err
    }

    dm.db.CreateIndex(&Submission{}, "pid")
    dm.db.CreateIndex(&Submission{}, "cid")
    dm.db.CreateIndex(&Submission{}, "iid")
    dm.db.CreateIndex(&Submission{}, "status")

    return nil
}

func (dm *DatabaseManager) SubmissionNew(pid, cid, iid, lang int64, code string) (int64, error) {
    sm := Submission {
        Pid: pid,
        Cid: cid,
        Iid: iid,
        Lang: lang,
        Time: 0,
        Mem: 0,
        Score: 0,
        SubmitTime: time.Now().Unix(),
        Status: int64(InQueue),
        Prog: 0,
    }

    id, err := dm.db.Insert(&sm)

    if err != nil {
        return 0, err
    }

    err = os.MkdirAll(SubmissionDir + strconv.FormatInt(id, 10), 0644)

    if err != nil {
        //SubmissionDelete

        return 0, err
    }

    fp, err := os.OpenFile(SubmissionDir + strconv.FormatInt(id, 10) + "/msg", os.O_WRONLY | os.O_CREATE, 0644)

    if err != nil {
        return 0, err
    }

    fp.Close()

    fp, err = os.OpenFile(SubmissionDir + strconv.FormatInt(id, 10) + "/case", os.O_WRONLY | os.O_CREATE, 0644)

    if err != nil {
        return 0, err
    }

    fp.Close()

    fp, err = os.OpenFile(SubmissionDir + strconv.FormatInt(id, 10) + "/code", os.O_WRONLY | os.O_CREATE, 0644)

    if err != nil {
        return 0, err
    }

    defer fp.Close()

    _, err = fp.Write([]byte(code))

    if err != nil {
        return 0, err
    }

    return id, nil
}

func (dm *DatabaseManager) SubmissionRemove(sid int64) error {
    sm := Submission{Sid: sid}
    _, err := dm.db.Delete(&sm)

    if err != nil {
        return err
    }

    fm1, _ := FileManager.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/case", os.O_RDONLY, true)
    fm2, _ := FileManager.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/msg", os.O_RDONLY, true)

    defer func () {
        fm1.Close()
        fm2.Close()
    }()

    return os.RemoveAll(SubmissionDir + strconv.FormatInt(sid, 10))
}

func (dm *DatabaseManager) SubmissionFind(sid int64) (*Submission, error) {
    var results []Submission

    err := dm.db.Select(&results, dm.db.Where("sid", "=", sid))

    if err != nil {
        return nil, err
    }

    if len(results) == 0 {
        return nil, errors.New("Unknown submission")
    }

    return &results[0], nil
}

func (dm *DatabaseManager) SubmissionUpdate(sid, time, mem int64, status SubmissionStatus, fin, all int) error {
    sm, err := dm.SubmissionFind(sid)

    if err != nil {
        return err
    }

    sm.Time = time
    sm.Mem = mem
    sm.Status = int64(status)
    sm.Prog = uint64(fin) << 32 & uint64(all)

    _, err = dm.db.Update(&sm)

    return err
}

func (dm *DatabaseManager) SubmissionGetCode(sid int64) (*string, error) {
    fp, err := os.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/code", os.O_RDONLY, 0644)

    if err != nil {
        return nil, err
    }

    defer fp.Close()

    b, err := ioutil.ReadAll(fp)

    if err != nil {
        return nil, err
    }

    str := string(b)

    return &str, nil
}

func (dm *DatabaseManager) SubmissionGetMsg(sid int64) string {
    fm, err := FileManager.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/msg", os.O_RDONLY, false)

    if err != nil {
        return ""
    }

    defer fm.Close()

    b, err := ioutil.ReadAll(fm)

    if err != nil {
        return ""
    }

    return string(b)
}

func (dm *DatabaseManager) SubmissionSetMsg(sid int64, msg string) error {
    fm, err := FileManager.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/msg", os.O_WRONLY | os.O_CREATE, true)

    if err != nil {
        return err
    }

    defer fm.Close()

    _, err = fm.Write([]byte(msg))

    return err
}

func (dm *DatabaseManager) SubmissionGetCase(sid int64) (*map[int64]SubmissionStatus, error) {
    fm, err := FileManager.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/case", os.O_RDONLY, false)

    if err != nil {
        return nil, err
    }

    defer fm.Close()

    b, err := ioutil.ReadAll(fm)

    if err != nil {
        return nil, err
    }

    var ss map[int64]SubmissionStatus
    err = json.Unmarshal(b, &ss)

    if err != nil {
        return nil, err
    }

    return &ss, nil
}

func (dm *DatabaseManager) SubmissionSetCase(sid int64, ss map[int64]SubmissionStatus) error {
    fm, err := FileManager.OpenFile(SubmissionDir + strconv.FormatInt(sid, 10) + "/case", os.O_RDONLY, false)

    if err != nil {
        return err
    }

    defer fm.Close()

    b, err := json.Marshal(ss)

    if err != nil {
        return err
    }

    _, err = fm.Write(b)

    return err
}<|endoftext|>