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

// This file implements the host side of CGI (being the webserver
// parent process).

// Package cgi implements CGI (Common Gateway Interface) as specified
// in RFC 3875.
//
// Note that using CGI means starting a new process to handle each
// request, which is typically less efficient than using a
// long-running server.  This package is intended primarily for
// compatibility with existing systems.
package cgi

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"syscall"
)

func GetFd(l interface{}) uintptr {
	v := reflect.ValueOf(l).Elem().FieldByName("fd")
	return uintptr(v.Int())
}

var trailingPort = regexp.MustCompile(`:([0-9]+)$`)

var osDefaultInheritEnv = map[string][]string{
	"darwin":  {"DYLD_LIBRARY_PATH"},
	"freebsd": {"LD_LIBRARY_PATH"},
	"hpux":    {"LD_LIBRARY_PATH", "SHLIB_PATH"},
	"irix":    {"LD_LIBRARY_PATH", "LD_LIBRARYN32_PATH", "LD_LIBRARY64_PATH"},
	"linux":   {"LD_LIBRARY_PATH"},
	"openbsd": {"LD_LIBRARY_PATH"},
	"solaris": {"LD_LIBRARY_PATH", "LD_LIBRARY_PATH_32", "LD_LIBRARY_PATH_64"},
	"windows": {"SystemRoot", "COMSPEC", "PATHEXT", "WINDIR"},
}

// Handler runs an executable in a subprocess with a CGI environment.
type Handler struct {
	Path string // path to the CGI executable
	Root string // root URI prefix of handler or empty for "/"

	// Dir specifies the CGI executable's working directory.
	// If Dir is empty, the base directory of Path is used.
	// If Path has no base directory, the current working
	// directory is used.
	Dir string

	Env        []string    // extra environment variables to set, if any, as "key=value"
	InheritEnv []string    // environment variables to inherit from host, as "key"
	Logger     *log.Logger // optional log for errors or nil to use log.Print
	Args       []string    // optional arguments to pass to child process
	Stderr     io.Writer   // optional place to send stderr of process to (default: os.Stderr)

	// PathLocationHandler specifies the root http Handler that
	// should handle internal redirects when the CGI process
	// returns a Location header value starting with a "/", as
	// specified in RFC 3875 § 6.3.2. This will likely be
	// http.DefaultServeMux.
	//
	// If nil, a CGI response with a local URI path is instead sent
	// back to the client and not redirected internally.
	PathLocationHandler http.Handler
}

// removeLeadingDuplicates remove leading duplicate in environments.
// It's possible to override environment like following.
//    cgi.Handler{
//      ...
//      Env: []string{"SCRIPT_FILENAME=foo.php"},
//    }
func removeLeadingDuplicates(env []string) (ret []string) {
	n := len(env)
	for i := 0; i < n; i++ {
		e := env[i]
		s := strings.SplitN(e, "=", 2)[0]
		found := false
		for j := i + 1; j < n; j++ {
			if s == strings.SplitN(env[j], "=", 2)[0] {
				found = true
				break
			}
		}
		if !found {
			ret = append(ret, e)
		}
	}
	return
}

func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	root := h.Root
	if root == "" {
		root = "/"
	}

	if len(req.TransferEncoding) > 0 && req.TransferEncoding[0] == "chunked" {
		rw.WriteHeader(http.StatusBadRequest)
		rw.Write([]byte("Chunked request bodies are not supported by CGI."))
		return
	}

	pathInfo := req.URL.Path
	if root != "/" && strings.HasPrefix(pathInfo, root) {
		pathInfo = pathInfo[len(root):]
	}

	port := "80"
	if matches := trailingPort.FindStringSubmatch(req.Host); len(matches) != 0 {
		port = matches[1]
	}

	env := []string{
		"SERVER_SOFTWARE=go",
		"SERVER_NAME=" + req.Host,
		"SERVER_PROTOCOL=HTTP/1.1",
		"HTTP_HOST=" + req.Host,
		"GATEWAY_INTERFACE=CGI/1.1",
		"REQUEST_METHOD=" + req.Method,
		"QUERY_STRING=" + req.URL.RawQuery,
		"REQUEST_URI=" + req.URL.RequestURI(),
		"PATH_INFO=" + pathInfo,
		"SCRIPT_NAME=" + root,
		"SCRIPT_FILENAME=" + h.Path,
		"REMOTE_ADDR=" + req.RemoteAddr,
		"REMOTE_HOST=" + req.RemoteAddr,
		"SERVER_PORT=" + port,
	}

	if req.TLS != nil {
		env = append(env, "HTTPS=on")
	}

	for k, v := range req.Header {
		k = strings.Map(upperCaseAndUnderscore, k)
		joinStr := ", "
		if k == "COOKIE" {
			joinStr = "; "
		}
		env = append(env, "HTTP_"+k+"="+strings.Join(v, joinStr))
	}

	if req.ContentLength > 0 {
		env = append(env, fmt.Sprintf("CONTENT_LENGTH=%d", req.ContentLength))
	}
	if ctype := req.Header.Get("Content-Type"); ctype != "" {
		env = append(env, "CONTENT_TYPE="+ctype)
	}

	if h.Env != nil {
		env = append(env, h.Env...)
	}

	envPath := os.Getenv("PATH")
	if envPath == "" {
		envPath = "/bin:/usr/bin:/usr/ucb:/usr/bsd:/usr/local/bin"
	}
	env = append(env, "PATH="+envPath)

	for _, e := range h.InheritEnv {
		if v := os.Getenv(e); v != "" {
			env = append(env, e+"="+v)
		}
	}

	for _, e := range osDefaultInheritEnv[runtime.GOOS] {
		if v := os.Getenv(e); v != "" {
			env = append(env, e+"="+v)
		}
	}

	env = removeLeadingDuplicates(env)

	var cwd, path string
	if h.Dir != "" {
		path = h.Path
		cwd = h.Dir
	} else {
		cwd, path = filepath.Split(h.Path)
	}
	if cwd == "" {
		cwd = "."
	}

	internalError := func(err error) {
		rw.WriteHeader(http.StatusInternalServerError)
		h.printf("CGI error: %v", err)
	}

	stderr := h.Stderr
	if stderr == nil {
		stderr = os.Stderr
	}

	cmd := &exec.Cmd{
		Path:   path,
		Args:   append([]string{h.Path}, h.Args...),
		Dir:    cwd,
		Env:    env,
		Stderr: stderr,
	}
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setsid: true,
	}
	if req.ContentLength != 0 {
		cmd.Stdin = req.Body
	}
	stdoutRead, err := cmd.StdoutPipe()
	if err != nil {
		internalError(err)
		return
	}

	err = cmd.Start()
	if err != nil {
		internalError(err)
		return
	}
	if hook := testHookStartProcess; hook != nil {
		hook(cmd.Process)
	}
	defer cmd.Wait()
	defer stdoutRead.Close()

	linebody := bufio.NewReaderSize(stdoutRead, 1024)
	headers := make(http.Header)
	statusCode := 0
	headerLines := 0
	sawBlankLine := false
	for {
		line, isPrefix, err := linebody.ReadLine()
		if isPrefix {
			rw.WriteHeader(http.StatusInternalServerError)
			h.printf("cgi: long header line from subprocess.")
			return
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			rw.WriteHeader(http.StatusInternalServerError)
			h.printf("cgi: error reading headers: %v", err)
			return
		}
		if len(line) == 0 {
			sawBlankLine = true
			break
		}
		headerLines++
		parts := strings.SplitN(string(line), ":", 2)
		if len(parts) < 2 {
			h.printf("cgi: bogus header line: %s", string(line))
			continue
		}
		header, val := parts[0], parts[1]
		header = strings.TrimSpace(header)
		val = strings.TrimSpace(val)
		switch {
		case header == "Status":
			if len(val) < 3 {
				h.printf("cgi: bogus status (short): %q", val)
				return
			}
			code, err := strconv.Atoi(val[0:3])
			if err != nil {
				h.printf("cgi: bogus status: %q", val)
				h.printf("cgi: line was %q", line)
				return
			}
			statusCode = code
		default:
			headers.Add(header, val)
		}
	}
	if headerLines == 0 || !sawBlankLine {
		rw.WriteHeader(http.StatusInternalServerError)
		h.printf("cgi: no headers")
		return
	}

	if loc := headers.Get("Location"); loc != "" {
		if strings.HasPrefix(loc, "/") && h.PathLocationHandler != nil {
			h.handleInternalRedirect(rw, req, loc)
			return
		}
		if statusCode == 0 {
			statusCode = http.StatusFound
		}
	}

	if statusCode == 0 && headers.Get("Content-Type") == "" {
		rw.WriteHeader(http.StatusInternalServerError)
		h.printf("cgi: missing required Content-Type in headers")
		return
	}

	if statusCode == 0 {
		statusCode = http.StatusOK
	}

	// Copy headers to rw's headers, after we've decided not to
	// go into handleInternalRedirect, which won't want its rw
	// headers to have been touched.
	for k, vv := range headers {
		for _, v := range vv {
			rw.Header().Add(k, v)
		}
	}

	rw.WriteHeader(statusCode)

	log.Println("Here in CGI")
	go func() {
		// log.Printf("Waiting for closenotify ", rw)
		<-rw.(http.CloseNotifier).CloseNotify()
		log.Println("CloseNotify() ", err)
		// stdoutRead.Close()
		// cmd.Process.Signal(os.Kill)
		// fd := GetFd(stdoutRead)
		// log.Printf("Arg = %T %v", stdoutRead, stdoutRead, fd)
		// syscall.Close(int(fd))
		// log.Println("Closed!")
		pgrp := &os.Process{Pid: -cmd.Process.Pid}
		err := pgrp.Signal(os.Kill)
		log.Println("SIGNAL ERROR:", err)
		// pgrp
		// pgrp.
		// os.Signal
		// cmd.Process.Pid
		// cmd.Stdin.Clos
	}()

	_, err = io.Copy(rw, linebody)
	if err != nil {
		h.printf("cgi: copy error: %v", err)
		// And kill the child CGI process so we don't hang on
		// the deferred cmd.Wait above if the error was just
		// the client (rw) going away. If it was a read error
		// (because the child died itself), then the extra
		// kill of an already-dead process is harmless (the PID
		// won't be reused until the Wait above).
		cmd.Process.Kill()
	}
	log.Println("Closed")
}

func (h *Handler) printf(format string, v ...interface{}) {
	if h.Logger != nil {
		h.Logger.Printf(format, v...)
	} else {
		log.Printf(format, v...)
	}
}

func (h *Handler) handleInternalRedirect(rw http.ResponseWriter, req *http.Request, path string) {
	url, err := req.URL.Parse(path)
	if err != nil {
		rw.WriteHeader(http.StatusInternalServerError)
		h.printf("cgi: error resolving local URI path %q: %v", path, err)
		return
	}
	// TODO: RFC 3875 isn't clear if only GET is supported, but it
	// suggests so: "Note that any message-body attached to the
	// request (such as for a POST request) may not be available
	// to the resource that is the target of the redirect."  We
	// should do some tests against Apache to see how it handles
	// POST, HEAD, etc. Does the internal redirect get the same
	// method or just GET? What about incoming headers?
	// (e.g. Cookies) Which headers, if any, are copied into the
	// second request?
	newReq := &http.Request{
		Method:     "GET",
		URL:        url,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Header:     make(http.Header),
		Host:       url.Host,
		RemoteAddr: req.RemoteAddr,
		TLS:        req.TLS,
	}
	h.PathLocationHandler.ServeHTTP(rw, newReq)
}

func upperCaseAndUnderscore(r rune) rune {
	switch {
	case r >= 'a' && r <= 'z':
		return r - ('a' - 'A')
	case r == '-':
		return '_'
	case r == '=':
		// Maybe not part of the CGI 'spec' but would mess up
		// the environment in any case, as Go represents the
		// environment as a slice of "key=value" strings.
		return '_'
	}
	// TODO: other transformations in spec or practice?
	return r
}

var testHookStartProcess func(*os.Process) // nil except for some tests
<commit_msg>Tidy up CGI handler<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file implements the host side of CGI (being the webserver
// parent process).

// Package cgi implements CGI (Common Gateway Interface) as specified
// in RFC 3875.
//
// Note that using CGI means starting a new process to handle each
// request, which is typically less efficient than using a
// long-running server.  This package is intended primarily for
// compatibility with existing systems.
package cgi

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"syscall"
)

func GetFd(l interface{}) uintptr {
	v := reflect.ValueOf(l).Elem().FieldByName("fd")
	return uintptr(v.Int())
}

var trailingPort = regexp.MustCompile(`:([0-9]+)$`)

var osDefaultInheritEnv = map[string][]string{
	"darwin":  {"DYLD_LIBRARY_PATH"},
	"freebsd": {"LD_LIBRARY_PATH"},
	"hpux":    {"LD_LIBRARY_PATH", "SHLIB_PATH"},
	"irix":    {"LD_LIBRARY_PATH", "LD_LIBRARYN32_PATH", "LD_LIBRARY64_PATH"},
	"linux":   {"LD_LIBRARY_PATH"},
	"openbsd": {"LD_LIBRARY_PATH"},
	"solaris": {"LD_LIBRARY_PATH", "LD_LIBRARY_PATH_32", "LD_LIBRARY_PATH_64"},
	"windows": {"SystemRoot", "COMSPEC", "PATHEXT", "WINDIR"},
}

// Handler runs an executable in a subprocess with a CGI environment.
type Handler struct {
	Path string // path to the CGI executable
	Root string // root URI prefix of handler or empty for "/"

	// Dir specifies the CGI executable's working directory.
	// If Dir is empty, the base directory of Path is used.
	// If Path has no base directory, the current working
	// directory is used.
	Dir string

	Env        []string    // extra environment variables to set, if any, as "key=value"
	InheritEnv []string    // environment variables to inherit from host, as "key"
	Logger     *log.Logger // optional log for errors or nil to use log.Print
	Args       []string    // optional arguments to pass to child process
	Stderr     io.Writer   // optional place to send stderr of process to (default: os.Stderr)

	// PathLocationHandler specifies the root http Handler that
	// should handle internal redirects when the CGI process
	// returns a Location header value starting with a "/", as
	// specified in RFC 3875 § 6.3.2. This will likely be
	// http.DefaultServeMux.
	//
	// If nil, a CGI response with a local URI path is instead sent
	// back to the client and not redirected internally.
	PathLocationHandler http.Handler
}

// removeLeadingDuplicates remove leading duplicate in environments.
// It's possible to override environment like following.
//    cgi.Handler{
//      ...
//      Env: []string{"SCRIPT_FILENAME=foo.php"},
//    }
func removeLeadingDuplicates(env []string) (ret []string) {
	n := len(env)
	for i := 0; i < n; i++ {
		e := env[i]
		s := strings.SplitN(e, "=", 2)[0]
		found := false
		for j := i + 1; j < n; j++ {
			if s == strings.SplitN(env[j], "=", 2)[0] {
				found = true
				break
			}
		}
		if !found {
			ret = append(ret, e)
		}
	}
	return
}

func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	root := h.Root
	if root == "" {
		root = "/"
	}

	if len(req.TransferEncoding) > 0 && req.TransferEncoding[0] == "chunked" {
		rw.WriteHeader(http.StatusBadRequest)
		rw.Write([]byte("Chunked request bodies are not supported by CGI."))
		return
	}

	pathInfo := req.URL.Path
	if root != "/" && strings.HasPrefix(pathInfo, root) {
		pathInfo = pathInfo[len(root):]
	}

	port := "80"
	if matches := trailingPort.FindStringSubmatch(req.Host); len(matches) != 0 {
		port = matches[1]
	}

	env := []string{
		"SERVER_SOFTWARE=go",
		"SERVER_NAME=" + req.Host,
		"SERVER_PROTOCOL=HTTP/1.1",
		"HTTP_HOST=" + req.Host,
		"GATEWAY_INTERFACE=CGI/1.1",
		"REQUEST_METHOD=" + req.Method,
		"QUERY_STRING=" + req.URL.RawQuery,
		"REQUEST_URI=" + req.URL.RequestURI(),
		"PATH_INFO=" + pathInfo,
		"SCRIPT_NAME=" + root,
		"SCRIPT_FILENAME=" + h.Path,
		"REMOTE_ADDR=" + req.RemoteAddr,
		"REMOTE_HOST=" + req.RemoteAddr,
		"SERVER_PORT=" + port,
	}

	if req.TLS != nil {
		env = append(env, "HTTPS=on")
	}

	for k, v := range req.Header {
		k = strings.Map(upperCaseAndUnderscore, k)
		joinStr := ", "
		if k == "COOKIE" {
			joinStr = "; "
		}
		env = append(env, "HTTP_"+k+"="+strings.Join(v, joinStr))
	}

	if req.ContentLength > 0 {
		env = append(env, fmt.Sprintf("CONTENT_LENGTH=%d", req.ContentLength))
	}
	if ctype := req.Header.Get("Content-Type"); ctype != "" {
		env = append(env, "CONTENT_TYPE="+ctype)
	}

	if h.Env != nil {
		env = append(env, h.Env...)
	}

	envPath := os.Getenv("PATH")
	if envPath == "" {
		envPath = "/bin:/usr/bin:/usr/ucb:/usr/bsd:/usr/local/bin"
	}
	env = append(env, "PATH="+envPath)

	for _, e := range h.InheritEnv {
		if v := os.Getenv(e); v != "" {
			env = append(env, e+"="+v)
		}
	}

	for _, e := range osDefaultInheritEnv[runtime.GOOS] {
		if v := os.Getenv(e); v != "" {
			env = append(env, e+"="+v)
		}
	}

	env = removeLeadingDuplicates(env)

	var cwd, path string
	if h.Dir != "" {
		path = h.Path
		cwd = h.Dir
	} else {
		cwd, path = filepath.Split(h.Path)
	}
	if cwd == "" {
		cwd = "."
	}

	internalError := func(err error) {
		rw.WriteHeader(http.StatusInternalServerError)
		h.printf("CGI error: %v", err)
	}

	stderr := h.Stderr
	if stderr == nil {
		stderr = os.Stderr
	}

	cmd := &exec.Cmd{
		Path:   path,
		Args:   append([]string{h.Path}, h.Args...),
		Dir:    cwd,
		Env:    env,
		Stderr: stderr,
	}
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setsid: true,
	}
	if req.ContentLength != 0 {
		cmd.Stdin = req.Body
	}
	stdoutRead, err := cmd.StdoutPipe()
	if err != nil {
		internalError(err)
		return
	}

	err = cmd.Start()
	if err != nil {
		internalError(err)
		return
	}
	if hook := testHookStartProcess; hook != nil {
		hook(cmd.Process)
	}
	defer cmd.Wait()
	defer stdoutRead.Close()

	completed := make(chan struct{})
	defer close(completed)

	go func() {
		<-rw.(http.CloseNotifier).CloseNotify()

		select {
		default:
			pgrp := &os.Process{Pid: -cmd.Process.Pid}
			err := pgrp.Signal(os.Kill)
			log.Println("Sent KILL to pgrp:", err)
		case <-completed:
			log.Println("Process completed")
		}
	}()

	linebody := bufio.NewReaderSize(stdoutRead, 1024)
	headers := make(http.Header)
	statusCode := 0
	headerLines := 0
	sawBlankLine := false
	for {
		line, isPrefix, err := linebody.ReadLine()
		if isPrefix {
			rw.WriteHeader(http.StatusInternalServerError)
			h.printf("cgi: long header line from subprocess.")
			return
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			rw.WriteHeader(http.StatusInternalServerError)
			h.printf("cgi: error reading headers: %v", err)
			return
		}
		if len(line) == 0 {
			sawBlankLine = true
			break
		}
		headerLines++
		parts := strings.SplitN(string(line), ":", 2)
		if len(parts) < 2 {
			h.printf("cgi: bogus header line: %s", string(line))
			continue
		}
		header, val := parts[0], parts[1]
		header = strings.TrimSpace(header)
		val = strings.TrimSpace(val)
		switch {
		case header == "Status":
			if len(val) < 3 {
				h.printf("cgi: bogus status (short): %q", val)
				return
			}
			code, err := strconv.Atoi(val[0:3])
			if err != nil {
				h.printf("cgi: bogus status: %q", val)
				h.printf("cgi: line was %q", line)
				return
			}
			statusCode = code
		default:
			headers.Add(header, val)
		}
	}
	if headerLines == 0 || !sawBlankLine {
		rw.WriteHeader(http.StatusInternalServerError)
		h.printf("cgi: no headers")
		return
	}

	if loc := headers.Get("Location"); loc != "" {
		if strings.HasPrefix(loc, "/") && h.PathLocationHandler != nil {
			h.handleInternalRedirect(rw, req, loc)
			return
		}
		if statusCode == 0 {
			statusCode = http.StatusFound
		}
	}

	if statusCode == 0 && headers.Get("Content-Type") == "" {
		rw.WriteHeader(http.StatusInternalServerError)
		h.printf("cgi: missing required Content-Type in headers")
		return
	}

	if statusCode == 0 {
		statusCode = http.StatusOK
	}

	// Copy headers to rw's headers, after we've decided not to
	// go into handleInternalRedirect, which won't want its rw
	// headers to have been touched.
	for k, vv := range headers {
		for _, v := range vv {
			rw.Header().Add(k, v)
		}
	}

	rw.WriteHeader(statusCode)

	_, err = io.Copy(rw, linebody)
	if err != nil {
		h.printf("cgi: copy error: %v", err)
		// And kill the child CGI process so we don't hang on
		// the deferred cmd.Wait above if the error was just
		// the client (rw) going away. If it was a read error
		// (because the child died itself), then the extra
		// kill of an already-dead process is harmless (the PID
		// won't be reused until the Wait above).
		cmd.Process.Kill()
	}
}

func (h *Handler) printf(format string, v ...interface{}) {
	if h.Logger != nil {
		h.Logger.Printf(format, v...)
	} else {
		log.Printf(format, v...)
	}
}

func (h *Handler) handleInternalRedirect(rw http.ResponseWriter, req *http.Request, path string) {
	url, err := req.URL.Parse(path)
	if err != nil {
		rw.WriteHeader(http.StatusInternalServerError)
		h.printf("cgi: error resolving local URI path %q: %v", path, err)
		return
	}
	// TODO: RFC 3875 isn't clear if only GET is supported, but it
	// suggests so: "Note that any message-body attached to the
	// request (such as for a POST request) may not be available
	// to the resource that is the target of the redirect."  We
	// should do some tests against Apache to see how it handles
	// POST, HEAD, etc. Does the internal redirect get the same
	// method or just GET? What about incoming headers?
	// (e.g. Cookies) Which headers, if any, are copied into the
	// second request?
	newReq := &http.Request{
		Method:     "GET",
		URL:        url,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Header:     make(http.Header),
		Host:       url.Host,
		RemoteAddr: req.RemoteAddr,
		TLS:        req.TLS,
	}
	h.PathLocationHandler.ServeHTTP(rw, newReq)
}

func upperCaseAndUnderscore(r rune) rune {
	switch {
	case r >= 'a' && r <= 'z':
		return r - ('a' - 'A')
	case r == '-':
		return '_'
	case r == '=':
		// Maybe not part of the CGI 'spec' but would mess up
		// the environment in any case, as Go represents the
		// environment as a slice of "key=value" strings.
		return '_'
	}
	// TODO: other transformations in spec or practice?
	return r
}

var testHookStartProcess func(*os.Process) // nil except for some tests
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
)

var (
	counter       = make(map[string]int)
	statsFilePath = "/var/lib/like.json"
)

func logIf(err error) {
	if err != nil {
		log.Println(err)
	}
}

func writeCounterData() {
	f, err := os.Create(statsFilePath)
	logIf(err)

	defer f.Close()

	data, err := json.Marshal(counter)
	logIf(err)

	if err == nil {
		_, err = f.Write(data)
		logIf(err)
	}
}

func buttonHandler(w http.ResponseWriter, r *http.Request) {
	object := r.URL.Path[1:]
	method := r.Method

	type Context struct {
		Object  string `json:"object"`
		Counter int    `json:"counter"`
	}

	// The first and only URL path defines the object name we use as an
	// identifier to retrieve the counter.
	if len(object) == 0 {
		http.Error(w, "No object given", http.StatusBadRequest)
		return
	}

	// In case this object is not yet registered, just add it to the list
	if _, ok := counter[object]; !ok {
		counter[object] = 0
	}

	// POST requests increae the counter by one
	if method == "POST" {
		counter[object] += 1
	}

	// Return the current counter for the given object name
	context := &Context{Object: object, Counter: counter[object]}
	response, err := json.Marshal(context)

	if err != nil {
		log.Println(err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}

	fmt.Fprint(w, string(response))
	writeCounterData()
}

func main() {
	log.Println("Starting server on port 8080...")

	// Load the data dump
	if _, err := os.Stat(statsFilePath); err == nil {
		content, err := ioutil.ReadFile(statsFilePath)
		logIf(err)

		err = json.Unmarshal(content, &counter)
		logIf(err)
	} else {
		log.Print("No stat file found, starting fresh")
	}

	http.HandleFunc("/", buttonHandler)
	http.ListenAndServe(":12100", nil)
}
<commit_msg>Only write stats file on changes<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
)

var (
	counter       = make(map[string]int)
	statsFilePath = "/var/lib/like.json"
)

func logIf(err error) {
	if err != nil {
		log.Println(err)
	}
}

func writeCounterData() {
	f, err := os.Create(statsFilePath)
	logIf(err)

	defer f.Close()

	data, err := json.Marshal(counter)
	logIf(err)

	if err == nil {
		_, err = f.Write(data)
		logIf(err)
	}
}

func buttonHandler(w http.ResponseWriter, r *http.Request) {
	object := r.URL.Path[1:]
	method := r.Method

	type Context struct {
		Object  string `json:"object"`
		Counter int    `json:"counter"`
	}

	// The first and only URL path defines the object name we use as an
	// identifier to retrieve the counter.
	if len(object) == 0 {
		http.Error(w, "No object given", http.StatusBadRequest)
		return
	}

	// In case this object is not yet registered, just add it to the list
	if _, ok := counter[object]; !ok {
		counter[object] = 0
	}

	// POST requests increae the counter by one
	if method == "POST" {
		counter[object] += 1
	}

	// Return the current counter for the given object name
	context := &Context{Object: object, Counter: counter[object]}
	response, err := json.Marshal(context)

	if err != nil {
		log.Println(err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}

	fmt.Fprint(w, string(response))

    if method == "POST" {
        writeCounterData()
    }
}

func main() {
	log.Println("Starting server on port 8080...")

	// Load the data dump
	if _, err := os.Stat(statsFilePath); err == nil {
		content, err := ioutil.ReadFile(statsFilePath)
		logIf(err)

		err = json.Unmarshal(content, &counter)
		logIf(err)
	} else {
		log.Print("No stat file found, starting fresh")
	}

	http.HandleFunc("/", buttonHandler)
	http.ListenAndServe(":12100", nil)
}
<|endoftext|>


<commit_before>package mysqltest

import (
	"bitbucket.org/liamstask/goose/lib/goose"
	mt "github.com/lestrrat/go-test-mysqld"
	"github.com/ziutek/mymysql/godrv"
    "fmt"
)

type TestMysqld struct {
    Server *mt.TestMysqld
}

func init() {
	godrv.Register("CREATE DATABASE IF NOT EXISTS test")
}

// New function launch new MySQL instance with default settings,
// and deploy SQLs in `dir` by `goose` automatically.
func New(dir string) (*TestMysqld, error) {
    mysql, err := mt.NewMysqld(nil)
	if err != nil {
		return nil, err
	}

    t := &TestMysqld{mysql}

    if err := t.deployDB(dir); err != nil {
        mysql.Stop()
        return nil, err
    }

	return &TestMysqld{mysql}, nil
}

func (t *TestMysqld) deployDB(dir string) error {
    conf, err := goose.NewDBConf(dir, "development", "")
	if err != nil {
        return err
	}
    conf.Driver.OpenStr = t.Dns()

	target, err := goose.GetMostRecentDBVersion(conf.MigrationsDir)
	if err != nil {
        return err
	}

	if err := goose.RunMigrations(conf, conf.MigrationsDir, target); err != nil {
		return err
	}

    return nil
}

// Dns returns connection info for mymysql driver
func (t *TestMysqld) Dns() string {
    return fmt.Sprintf("unix:%s*test/root/", t.Server.Socket())
}
<commit_msg>create test database manually if its not exists<commit_after>package mysqltest

import (
	"bitbucket.org/liamstask/goose/lib/goose"
	mt "github.com/lestrrat/go-test-mysqld"
    "fmt"

	"database/sql"
	_ "github.com/ziutek/mymysql/godrv"
)

type TestMysqld struct {
    Server *mt.TestMysqld
}

// New function launch new MySQL instance with default settings,
// and deploy SQLs in `dir` by `goose` automatically.
func New(dir string) (*TestMysqld, error) {
    mysql, err := mt.NewMysqld(nil)
	if err != nil {
		return nil, err
	}

	db, err := sql.Open("mymysql", fmt.Sprintf("unix:%s*mysql/root/", mysql.Socket()))
	if err != nil {
		return nil, err
	}
	defer db.Close()

	_, err = db.Exec("CREATE DATABASE IF NOT EXISTS test")
	if err != nil {
		return nil, err
	}

    t := &TestMysqld{mysql}

    if err := t.deployDB(dir); err != nil {
        mysql.Stop()
        return nil, err
    }

	return &TestMysqld{mysql}, nil
}

func (t *TestMysqld) deployDB(dir string) error {
    conf, err := goose.NewDBConf(dir, "development", "")
	if err != nil {
        return err
	}
    conf.Driver.OpenStr = t.Dns()

	target, err := goose.GetMostRecentDBVersion(conf.MigrationsDir)
	if err != nil {
        return err
	}

	if err := goose.RunMigrations(conf, conf.MigrationsDir, target); err != nil {
		return err
	}

    return nil
}

// Dns returns connection info for mymysql driver
func (t *TestMysqld) Dns() string {
    return fmt.Sprintf("unix:%s*test/root/", t.Server.Socket())
}
<|endoftext|>


<commit_before>/*
Package liner implements a simple command line editor, inspired by linenoise
(https://github.com/antirez/linenoise/). This package supports WIN32 in
addition to the xterm codes supported by everything else.
*/
package liner

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"unicode/utf8"
)

type action int

const (
	left action = iota
	right
	up
	down
	home
	end
	insert
	del
	pageUp
	pageDown
	f1
	f2
	f3
	f4
	f5
	f6
	f7
	f8
	f9
	f10
	f11
	f12
	unknown
)

type commonState struct {
	history   []string
	supported bool
}

// ReadHistory reads scrollback history from r. Returns the number of lines
// read, and any read error (except io.EOF).
func (s *State) ReadHistory(r io.Reader) (num int, err error) {
	in := bufio.NewReader(r)
	num = 0
	for {
		line, part, err := in.ReadLine()
		if err == io.EOF {
			break
		}
		if err != nil {
			return num, err
		}
		if part {
			return num, errors.New("Line too long")
		}
		if !utf8.Valid(line) {
			return num, errors.New("Invalid string")
		}
		num++
		s.history = append(s.history, string(line))
	}
	return num, nil
}

// WriteHistory writes scrollback history to w. Returns the number of lines
// successfully written, and any write error.
func (s *State) WriteHistory(w io.Writer) (num int, err error) {
	for _, item := range s.history {
		_, err := fmt.Fprintln(w, item)
		if err != nil {
			return num, err
		}
		num++
	}
	return num, nil
}

const (
	ctrlA = 1
	ctrlB = 2
	ctrlC = 3
	ctrlD = 4
	ctrlE = 5
	ctrlF = 6
	ctrlH = 8
	lf    = 10
	ctrlK = 11
	ctrlL = 12
	cr    = 13
	ctrlN = 14
	ctrlP = 16
	ctrlT = 20
	ctrlU = 21
	esc   = 27
	bs    = 127
)

const (
	beep = "\a"
)

func (s *State) refresh(prompt string, buf string, pos int) error {
	s.cursorPos(0)
	_, err := fmt.Print(prompt)
	if err != nil {
		return err
	}
	_, err = fmt.Print(buf)
	s.eraseLine()
	s.cursorPos(utf8.RuneCountInString(prompt) + pos)
	return err
}

// Prompt displays p, and then waits for user input. Prompt allows line editing
// if the terminal supports it.
func (s *State) Prompt(p string) (string, error) {
	if !s.supported {
		return s.promptUnsupported(p)
	}

	fmt.Print(p)
	line := make([]rune, 0)
	pos := 0
mainLoop:
	for {
		next, err := s.readNext()
		if err != nil {
			return "", err
		}
		switch v := next.(type) {
		case rune:
			switch v {
			case cr, lf:
				break mainLoop
			case ctrlA: // Start of line
				pos = 0
				s.refresh(p, string(line), pos)
			case ctrlE: // End of line
				pos = len(line)
				s.refresh(p, string(line), pos)
			case ctrlB: // left
				if pos > 0 {
					pos--
					s.refresh(p, string(line), pos)
				} else {
					fmt.Print(beep)
				}
			case ctrlF: // right
				if pos < len(line) {
					pos++
					s.refresh(p, string(line), pos)
				} else {
					fmt.Print(beep)
				}
			case ctrlD: // del
				if pos >= len(line) {
					fmt.Print(beep)
				} else {
					line = append(line[:pos], line[pos+1:]...)
					s.refresh(p, string(line), pos)
				}
			case ctrlK: // delete remainder of line
				if pos >= len(line) {
					fmt.Print(beep)
				} else {
					line = line[:pos]
					s.refresh(p, string(line), pos)
				}
			case ctrlT: // transpose prev rune with rune under cursor
				if len(line) < 2 || pos < 1 {
					fmt.Print(beep)
				} else {
					if pos == len(line) {
						pos--
					}
					line[pos-1], line[pos] = line[pos], line[pos-1]
					pos++
					s.refresh(p, string(line), pos)
				}
			case ctrlL: // clear screen
				s.eraseScreen()
				s.refresh(p, string(line), pos)
			case ctrlH, bs: // Backspace
				if pos <= 0 {
					fmt.Print(beep)
				} else {
					line = append(line[:pos-1], line[pos:]...)
					pos--
					s.refresh(p, string(line), pos)
				}
			case ctrlU: // Erase entire line
				line = line[:0]
				pos = 0
				s.refresh(p, string(line), pos)
			// Catch unhandled control codes (anything <= 31)
			case 0, 3, 7, 9, 14, 15, 16:
				fallthrough
			case 17, 18, 19, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31:
				fmt.Print(beep)
			default:
				if pos == len(line) {
					line = append(line, v)
					fmt.Printf("%c", v)
					pos++
				} else {
					line = append(line[:pos], append([]rune{v}, line[pos:]...)...)
					pos++
					s.refresh(p, string(line), pos)
				}
			}
		case action:
			switch v {
			case del:
				if pos >= len(line) {
					fmt.Print(beep)
				} else {
					line = append(line[:pos], line[pos+1:]...)
				}
			case left:
				if pos > 0 {
					pos--
				} else {
					fmt.Print(beep)
				}
			case right:
				if pos < len(line) {
					pos++
				} else {
					fmt.Print(beep)
				}
			}
			s.refresh(p, string(line), pos)
		}
	}
	return string(line), nil
}
<commit_msg>Add history support<commit_after>/*
Package liner implements a simple command line editor, inspired by linenoise
(https://github.com/antirez/linenoise/). This package supports WIN32 in
addition to the xterm codes supported by everything else.
*/
package liner

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"unicode/utf8"
)

// HistoryLimit is the maximum number of entries saved in the scrollback history.
const HistoryLimit = 100

type action int

const (
	left action = iota
	right
	up
	down
	home
	end
	insert
	del
	pageUp
	pageDown
	f1
	f2
	f3
	f4
	f5
	f6
	f7
	f8
	f9
	f10
	f11
	f12
	unknown
)

type commonState struct {
	history   []string
	supported bool
}

// ReadHistory reads scrollback history from r. Returns the number of lines
// read, and any read error (except io.EOF).
func (s *State) ReadHistory(r io.Reader) (num int, err error) {
	in := bufio.NewReader(r)
	num = 0
	for {
		line, part, err := in.ReadLine()
		if err == io.EOF {
			break
		}
		if err != nil {
			return num, err
		}
		if part {
			return num, errors.New("Line too long")
		}
		if !utf8.Valid(line) {
			return num, errors.New("Invalid string")
		}
		num++
		s.history = append(s.history, string(line))
		if len(s.history) > HistoryLimit {
			s.history = s.history[1:]
		}
	}
	return num, nil
}

// WriteHistory writes scrollback history to w. Returns the number of lines
// successfully written, and any write error.
func (s *State) WriteHistory(w io.Writer) (num int, err error) {
	for _, item := range s.history {
		_, err := fmt.Fprintln(w, item)
		if err != nil {
			return num, err
		}
		num++
	}
	return num, nil
}

// AppendHistory appends an entry to the scrollback history. AppendHistory
// should be called iff Prompt returns a valid command.
func (s *State) AppendHistory(item string) {
	if len(s.history) > 0 {
		if item == s.history[len(s.history)-1] {
			return
		}
	}
	s.history = append(s.history, item)
	if len(s.history) > HistoryLimit {
		s.history = s.history[1:]
	}
}

const (
	ctrlA = 1
	ctrlB = 2
	ctrlC = 3
	ctrlD = 4
	ctrlE = 5
	ctrlF = 6
	ctrlH = 8
	lf    = 10
	ctrlK = 11
	ctrlL = 12
	cr    = 13
	ctrlN = 14
	ctrlP = 16
	ctrlT = 20
	ctrlU = 21
	esc   = 27
	bs    = 127
)

const (
	beep = "\a"
)

func (s *State) refresh(prompt string, buf string, pos int) error {
	s.cursorPos(0)
	_, err := fmt.Print(prompt)
	if err != nil {
		return err
	}
	_, err = fmt.Print(buf)
	s.eraseLine()
	s.cursorPos(utf8.RuneCountInString(prompt) + pos)
	return err
}

// Prompt displays p, and then waits for user input. Prompt allows line editing
// if the terminal supports it.
func (s *State) Prompt(p string) (string, error) {
	if !s.supported {
		return s.promptUnsupported(p)
	}

	fmt.Print(p)
	line := make([]rune, 0)
	pos := 0
	historyPos := len(s.history)
	var historyEnd string

mainLoop:
	for {
		next, err := s.readNext()
		if err != nil {
			return "", err
		}
		switch v := next.(type) {
		case rune:
			switch v {
			case cr, lf:
				break mainLoop
			case ctrlA: // Start of line
				pos = 0
				s.refresh(p, string(line), pos)
			case ctrlE: // End of line
				pos = len(line)
				s.refresh(p, string(line), pos)
			case ctrlB: // left
				if pos > 0 {
					pos--
					s.refresh(p, string(line), pos)
				} else {
					fmt.Print(beep)
				}
			case ctrlF: // right
				if pos < len(line) {
					pos++
					s.refresh(p, string(line), pos)
				} else {
					fmt.Print(beep)
				}
			case ctrlD: // del
				if pos >= len(line) {
					fmt.Print(beep)
				} else {
					line = append(line[:pos], line[pos+1:]...)
					s.refresh(p, string(line), pos)
				}
			case ctrlK: // delete remainder of line
				if pos >= len(line) {
					fmt.Print(beep)
				} else {
					line = line[:pos]
					s.refresh(p, string(line), pos)
				}
			case ctrlP: // up
				if historyPos > 0 {
					if historyPos == len(s.history) {
						historyEnd = string(line)
					}
					historyPos--
					line = []rune(s.history[historyPos])
					pos = len(line)
					s.refresh(p, string(line), pos)
				} else {
					fmt.Print(beep)
				}
			case ctrlN: // down
				if historyPos < len(s.history) {
					historyPos++
					if historyPos == len(s.history) {
						line = []rune(historyEnd)
					} else {
						line = []rune(s.history[historyPos])
					}
					pos = len(line)
					s.refresh(p, string(line), pos)
				} else {
					fmt.Print(beep)
				}
			case ctrlT: // transpose prev rune with rune under cursor
				if len(line) < 2 || pos < 1 {
					fmt.Print(beep)
				} else {
					if pos == len(line) {
						pos--
					}
					line[pos-1], line[pos] = line[pos], line[pos-1]
					pos++
					s.refresh(p, string(line), pos)
				}
			case ctrlL: // clear screen
				s.eraseScreen()
				s.refresh(p, string(line), pos)
			case ctrlH, bs: // Backspace
				if pos <= 0 {
					fmt.Print(beep)
				} else {
					line = append(line[:pos-1], line[pos:]...)
					pos--
					s.refresh(p, string(line), pos)
				}
			case ctrlU: // Erase entire line
				line = line[:0]
				pos = 0
				s.refresh(p, string(line), pos)
			// Catch unhandled control codes (anything <= 31)
			case 0, 3, 7, 9, 15:
				fallthrough
			case 17, 18, 19, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31:
				fmt.Print(beep)
			default:
				if pos == len(line) {
					line = append(line, v)
					fmt.Printf("%c", v)
					pos++
				} else {
					line = append(line[:pos], append([]rune{v}, line[pos:]...)...)
					pos++
					s.refresh(p, string(line), pos)
				}
			}
		case action:
			switch v {
			case del:
				if pos >= len(line) {
					fmt.Print(beep)
				} else {
					line = append(line[:pos], line[pos+1:]...)
				}
			case left:
				if pos > 0 {
					pos--
				} else {
					fmt.Print(beep)
				}
			case right:
				if pos < len(line) {
					pos++
				} else {
					fmt.Print(beep)
				}
			case up:
				if historyPos > 0 {
					if historyPos == len(s.history) {
						historyEnd = string(line)
					}
					historyPos--
					line = []rune(s.history[historyPos])
					pos = len(line)
				} else {
					fmt.Print(beep)
				}
			case down:
				if historyPos < len(s.history) {
					historyPos++
					if historyPos == len(s.history) {
						line = []rune(historyEnd)
					} else {
						line = []rune(s.history[historyPos])
					}
					pos = len(line)
				} else {
					fmt.Print(beep)
				}
			}
			s.refresh(p, string(line), pos)
		}
	}
	return string(line), nil
}
<|endoftext|>


<commit_before>// mogul - distributed lock via mongoDB
//
// Copyright 2016 - Rene Kroon <kroon.r.w@gmail.com>
//

package mogul

import (
	"time"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

// Mutex is the object used for locking
type Mutex struct {
	doc        lock
	collection *mgo.Collection
}

type lock struct {
	Name         string    `bson:"_id"`
	User         string    `bson:"user"`
	ExpiresAtUtc time.Time `bson:"expires"`
}

// MutexCreator is a specialization of the Manager struct for locking
type MutexCreator interface {
	NewMutex(name string, user string) *Mutex
}

// Trylock will claim a lock if it is available. It also returns true when you already hold the lock.
// This extends the duration if you already hold the lock.
func (m *Mutex) TryLock(atMost time.Duration) (bool, error) {

	now := time.Now().UTC()
	until := now.Add(atMost)

	m.doc.ExpiresAtUtc = until

	selector := bson.M{"$or": []bson.M{m.identityClause(), bson.M{"$and": []bson.M{bson.M{"_id": m.doc.Name}, bson.M{"expires": bson.M{"$lt": now}}}}}}
	_, err := m.collection.Upsert(selector, &m.doc)

	found := 0
	if err == nil {
		found, err = m.collection.Find(m.identityClause()).Count()
	}

	return found > 0, err
}

// Unlock frees the lock, removing the corresponding record in the database
func (m *Mutex) Unlock() error {
	return m.collection.Remove(m.identityClause())
}

// IsExpired will return true when your lock time has expired
func (m *Mutex) IsExpired() bool {
	return m.doc.ExpiresAtUtc.Before(time.Now().UTC())
}

func (m *Mutex) identityClause() bson.M {
	return bson.M{"$and": []bson.M{bson.M{"_id": m.doc.Name}, bson.M{"user": m.doc.User}}}
}
<commit_msg>Fix race condition<commit_after>// mogul - distributed lock via mongoDB
//
// Copyright 2016 - Rene Kroon <kroon.r.w@gmail.com>
//

package mogul

import (
	"time"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"log"
)

// Mutex is the object used for locking
type Mutex struct {
	doc        lock
	collection *mgo.Collection
}

type lock struct {
	Name         string    `bson:"_id"`
	User         string    `bson:"user"`
	ExpiresAtUtc time.Time `bson:"expires"`
}

// MutexCreator is a specialization of the Manager struct for locking
type MutexCreator interface {
	NewMutex(name string, user string) *Mutex
}

// Trylock will claim a lock if it is available. It also returns true when you already hold the lock.
// This extends the duration if you already hold the lock.
func (m *Mutex) TryLock(atMost time.Duration) (bool, error) {

	now := time.Now().UTC()
	until := now.Add(atMost)

	m.doc.ExpiresAtUtc = until

	var result lock

	selector := bson.M{"$or": []bson.M{m.identityClause(), bson.M{"$and": []bson.M{bson.M{"_id": m.doc.Name}, bson.M{"expires": bson.M{"$lt": now}}}}}}

	_, err := m.collection.Find(selector).Apply(mgo.Change{
		Update:    &m.doc,
		ReturnNew: true,
	}, &result)
	if err == mgo.ErrNotFound {
		err = m.collection.Insert(&m.doc)
	}

	return err == nil, err
}

// Unlock frees the lock, removing the corresponding record in the database
func (m *Mutex) Unlock() error {
	return m.collection.Remove(m.identityClause())
}

// IsExpired will return true when your lock time has expired
func (m *Mutex) IsExpired() bool {
	return m.doc.ExpiresAtUtc.Before(time.Now().UTC())
}

func (m *Mutex) identityClause() bson.M {
	return bson.M{"$and": []bson.M{bson.M{"_id": m.doc.Name}, bson.M{"user": m.doc.User}}}
}
<|endoftext|>


<commit_before>package corgis

import (
	"encoding/json"
	"io"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"time"

	_ "github.com/jinzhu/gorm/dialects/postgres"
	"gopkg.in/pipe.v2"
)

func GetArguments(pid int) []string {
	if pid == 0 {
		return nil
	}
	filename := "/proc/" + strconv.Itoa(pid) + "/cmdline"
	p := pipe.Line(
		pipe.ReadFile(filename),
		pipe.Exec("strings", "-1"),
	)
	output, err := pipe.CombinedOutput(p)
	if err != nil {
		//fmt.Printf("error:[%v]\n", err)
	}
	return strings.Fields(string(output))
}
func timedSIGTERM(p *os.Process, d time.Duration) {
	//log.Printf("timedSIGTERM called\n")
	var _ = <-time.After(d)
	err := p.Signal(syscall.SIGTERM)
	//log.Printf("SIGTERM sent\n")
	if err != nil {
		log.Panicf("timedSIGTERM Panicking: %v\n", err)
	}
}

func TimedSIGTERM(p *os.Process, d time.Duration) {
	timedSIGTERM(p, d)
}

func DecodeStream(r io.Reader) []interface{} {
	decoder := json.NewDecoder(r)
	openToken, err := decoder.Token()
	if err != nil {
		log.Fatalf("openToken error: %v\n", err)
	}
	//fmt.Println("openToken:", openToken)
	var _ = openToken
	res := make([]interface{}, 0)
	for decoder.More() {
		var m interface{}
		err := decoder.Decode(&m)
		if err != nil {
			log.Fatalf("error decoding: %v\n", err)
		}
		res = append(res, m)
		//fmt.Printf("decoded: [%v]\n", m)
	}
	closeToken, err := decoder.Token()
	if err != nil {
		log.Fatalf("closeToken error: %v\n", err)
	}
	//fmt.Println("closeToken:", closeToken)
	var _ = closeToken
	return res
}

func (j *JobScheduler) InitCmd() {
	j.Cmd = exec.Command(j.Cmd.Path, j.Cmd.Args[1])
	j.OutBuf.Reset()
	j.ErrBuf.Reset()
	j.Cmd.Stdout = &j.OutBuf
	j.Cmd.Stderr = &j.ErrBuf
	//log.Printf("Initialized job\n")
}

// Curently uses for read&write latencies
func (j *JobScheduler) ExecStreaming() {
	err := j.Cmd.Start()
	if err != nil {
		log.Panicf("cmd start error: %v\n", err)
	}
	//log.Printf("job started\n")
	go timedSIGTERM(j.Cmd.Process, j.ExecPeriod)
	//log.Printf("waiting for job terminate\n")
	err = j.Cmd.Wait()
	if err != nil {
		log.Panicf("cmd wait error: %v\n", err)
	}
	if j.Cmd.ProcessState.Success() {
		j.Res = DecodeStream(&j.OutBuf)
		m := make(map[string]float64)
		for _, e := range j.Res {
			tastd, _ := e.(map[string]interface{})
			procName, _ := tastd["execname"].(string)
			if procName == "qemu-kvm" {
				// fmt.Printf("%v %v\n", i, tastd)
				tpid, _ := tastd["pid"].(float64)
				ipid := int(tpid)
				vmName := GetArguments(ipid)[2]
				ltc, _ := tastd["latency"].(float64)
				m[vmName] += ltc
			}
		}
		//fmt.Printf("%v %v\n", len(m), m)
		for i, e := range m {
			if j.LType == LRead {

				var storageType TiramisuStorage
				DB.First(&storageType, "vm_name = ?", i)

				vminfo := RawVMData{
					VMName:      i,
					LatencyRead: e / j.ExecPeriod.Seconds(),
					ISSSD:       storageType.CurrentPool == "SSD",
				}
				//fmt.Printf("%v\n", vminfo)
				DB.Where(RawVMData{VMName: i}).Assign(vminfo).FirstOrCreate(&vminfo)
			} else if j.LType == LWrite {

				var storageType TiramisuStorage
				DB.First(&storageType, "vm_name = ?", i)

				vminfo := RawVMData{
					VMName:       i,
					LatencyWrite: e / j.ExecPeriod.Seconds(),
					ISSSD:        storageType.CurrentPool == "SSD",
				}
				//fmt.Printf("%v\n", vminfo)
				DB.Where(RawVMData{VMName: i}).Assign(vminfo).FirstOrCreate(&vminfo)

			}
		}
		j.InitCmd()
		j.ExecStreaming()
	}

}

// Currently uses for IOPS
func (j *JobScheduler) ExecTimed() {
	err := j.Cmd.Start()
	if err != nil {
		log.Panicf("cmd start error: %v\n", err)
	}
	log.Printf("job started\n")
	go timedSIGTERM(j.Cmd.Process, j.ExecPeriod)
	log.Printf("waiting for job terminate\n")
	err = j.Cmd.Wait()
	if err != nil {
		log.Panicf("cmd wait error: %v\n", err)
	}
	if j.Cmd.ProcessState.Success() {
		j.Res = DecodeStream(&j.OutBuf)
		//fmt.Printf("-->%v\n", j.Res)
		for _, e := range j.Res {
			tmp, _ := e.(map[string]interface{})
			s_pid, _ := tmp["pid"].(string)
			var e_pid int
			if len(s_pid) != 0 {
				e_pid, err = strconv.Atoi(s_pid)
				if err != nil {
					log.Fatalf("error convert atoi: %v\n", err)
				}
			}
			argList := GetArguments(e_pid)
			if len(argList) != 0 {
				if argList[0] == "/usr/libexec/qemu-kvm" {
					//fmt.Printf("%v %v %v\n", i, tmp, argList[2])
					ioread, _ := tmp["read"].(float64)
					iowrite, _ := tmp["write"].(float64)

					var storageType TiramisuStorage
					DB.First(&storageType, "vm_name = ?", argList[2])
					// fmt.Printf("\n\n---->[%v]\n\n", storageType)

					vminfo := RawVMData{
						VMName:    argList[2],
						IOPSRead:  ioread / j.ExecPeriod.Seconds(),
						IOPSWrite: iowrite / j.ExecPeriod.Seconds(),
						ISSSD:     storageType.CurrentPool == "SSD",
					}
					//fmt.Printf("\n\n---->[%v]\n\n", vminfo)
					//fmt.Printf("\n[%v]\n", vminfo)
					DB.Where(RawVMData{VMName: argList[2]}).Assign(vminfo).FirstOrCreate(&vminfo)
				}
			}
		}
		j.InitCmd()
		j.ExecTimed()
	}
}

func (j *JobScheduler) Execute() {
	if j.Type == Streaming {
		j.ExecStreaming()

	} else if j.Type == Timed {
		j.ExecTimed()
	}
}

func AssignAverage() {
	log.Printf("AssignAverage Invoked\n")
	var HDDVMList []RawVMData
	var SSDVMList []RawVMData
	DB.Model(&RawVMData{}).Where("isssd = ?", true).Find(&SSDVMList)
	DB.Model(&RawVMData{}).Where("isssd = ?", false).Find(&HDDVMList)

	var avgHDDLatency float64
	var avgSSDLatency float64
	var avgSSDIOPS float64
	var avgHDDIOPS float64

	for _, e := range HDDVMList {
		avgHDDIOPS += e.IOPSRead + e.IOPSWrite
		avgHDDLatency += e.LatencyRead + e.LatencyWrite
	}
	for _, e := range SSDVMList {
		avgSSDIOPS += e.IOPSRead + e.IOPSWrite
		avgSSDLatency += e.LatencyRead + e.LatencyWrite
	}
	avgHDDIOPS = (0.5 * avgHDDIOPS) / float64(len(HDDVMList))
	avgHDDLatency = (0.5 * avgHDDLatency) / float64(len(HDDVMList))
	avgSSDIOPS = (0.5 * avgSSDIOPS) / float64(len(SSDVMList))
	avgSSDLatency = (0.5 * avgSSDLatency) / float64(len(SSDVMList))

	for _, e := range HDDVMList {
		tiramisu_state := TiramisuState{
			IOPS:        (e.IOPSRead + e.IOPSWrite) / 2,
			Latency:     (e.LatencyRead + e.LatencyWrite) / 2,
			IOPS_HDD:    (e.IOPSRead + e.IOPSWrite) / 2,
			Latency_HDD: (e.LatencyRead + e.LatencyWrite) / 2,
			IOPS_SSD:    avgSSDIOPS,
			Latency_SSD: avgSSDLatency,
			Name:        e.VMName,
		}
		DB.Model(&TiramisuState{}).Where("vm_name = ?", e.VMName).Assign(tiramisu_state).FirstOrCreate(&tiramisu_state)
	}
	for _, e := range SSDVMList {
		tiramisu_state := TiramisuState{
			IOPS:        (e.IOPSRead + e.IOPSWrite) / 2,
			Latency:     (e.LatencyRead + e.LatencyWrite) / 2,
			IOPS_HDD:    avgHDDIOPS,
			Latency_HDD: avgHDDLatency,
			IOPS_SSD:    (e.IOPSRead + e.IOPSWrite) / 2,
			Latency_SSD: (e.LatencyRead + e.LatencyWrite) / 2,
			Name:        e.VMName,
		}
		DB.Model(&TiramisuState{}).Where("vm_name = ?", e.VMName).Assign(tiramisu_state).FirstOrCreate(&tiramisu_state)
	}
}
<commit_msg>mark for refactor<commit_after>package corgis

import (
	"encoding/json"
	"io"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"time"

	_ "github.com/jinzhu/gorm/dialects/postgres"
	"gopkg.in/pipe.v2"
)

func GetArguments(pid int) []string {
	if pid == 0 {
		return nil
	}
	filename := "/proc/" + strconv.Itoa(pid) + "/cmdline"
	p := pipe.Line(
		pipe.ReadFile(filename),
		pipe.Exec("strings", "-1"),
	)
	output, err := pipe.CombinedOutput(p)
	if err != nil {
		//fmt.Printf("error:[%v]\n", err)
	}
	return strings.Fields(string(output))
}
func timedSIGTERM(p *os.Process, d time.Duration) {
	//log.Printf("timedSIGTERM called\n")
	var _ = <-time.After(d)
	err := p.Signal(syscall.SIGTERM)
	//log.Printf("SIGTERM sent\n")
	if err != nil {
		log.Panicf("timedSIGTERM Panicking: %v\n", err)
	}
}

func TimedSIGTERM(p *os.Process, d time.Duration) {
	timedSIGTERM(p, d)
}

func DecodeStream(r io.Reader) []interface{} {
	decoder := json.NewDecoder(r)
	openToken, err := decoder.Token()
	if err != nil {
		log.Fatalf("openToken error: %v\n", err)
	}
	//fmt.Println("openToken:", openToken)
	var _ = openToken
	res := make([]interface{}, 0)
	for decoder.More() {
		var m interface{}
		err := decoder.Decode(&m)
		if err != nil {
			log.Fatalf("error decoding: %v\n", err)
		}
		res = append(res, m)
		//fmt.Printf("decoded: [%v]\n", m)
	}
	closeToken, err := decoder.Token()
	if err != nil {
		log.Fatalf("closeToken error: %v\n", err)
	}
	//fmt.Println("closeToken:", closeToken)
	var _ = closeToken
	return res
}

func (j *JobScheduler) InitCmd() {
	j.Cmd = exec.Command(j.Cmd.Path, j.Cmd.Args[1])
	j.OutBuf.Reset()
	j.ErrBuf.Reset()
	j.Cmd.Stdout = &j.OutBuf
	j.Cmd.Stderr = &j.ErrBuf
	//log.Printf("Initialized job\n")
}

// Curently uses for read&write latencies
func (j *JobScheduler) ExecStreaming() {
	err := j.Cmd.Start()
	if err != nil {
		log.Panicf("cmd start error: %v\n", err)
	}
	//log.Printf("job started\n")
	go timedSIGTERM(j.Cmd.Process, j.ExecPeriod)
	//log.Printf("waiting for job terminate\n")
	err = j.Cmd.Wait()
	if err != nil {
		log.Panicf("cmd wait error: %v\n", err)
	}
	if j.Cmd.ProcessState.Success() {
		j.Res = DecodeStream(&j.OutBuf)
		m := make(map[string]float64)
		for _, e := range j.Res {
			tastd, _ := e.(map[string]interface{})
			procName, _ := tastd["execname"].(string)
			if procName == "qemu-kvm" {
				// fmt.Printf("%v %v\n", i, tastd)
				tpid, _ := tastd["pid"].(float64)
				ipid := int(tpid)
				vmName := GetArguments(ipid)[2]
				ltc, _ := tastd["latency"].(float64)
				m[vmName] += ltc
			}
		}
		//fmt.Printf("%v %v\n", len(m), m)
		for i, e := range m {
			if j.LType == LRead {

				var storageType TiramisuStorage
				DB.First(&storageType, "vm_name = ?", i)

				vminfo := RawVMData{
					VMName:      i,
					LatencyRead: e / j.ExecPeriod.Seconds(),
					ISSSD:       storageType.CurrentPool == "SSD",
				}
				//fmt.Printf("%v\n", vminfo)
				DB.Where(RawVMData{VMName: i}).Assign(vminfo).FirstOrCreate(&vminfo)
			} else if j.LType == LWrite {

				var storageType TiramisuStorage
				DB.First(&storageType, "vm_name = ?", i)

				vminfo := RawVMData{
					VMName:       i,
					LatencyWrite: e / j.ExecPeriod.Seconds(),
					ISSSD:        storageType.CurrentPool == "SSD",
				}
				//fmt.Printf("%v\n", vminfo)
				DB.Where(RawVMData{VMName: i}).Assign(vminfo).FirstOrCreate(&vminfo)

			}
		}
		j.InitCmd()
		j.ExecStreaming()
	}

}

// Currently uses for IOPS
func (j *JobScheduler) ExecTimed() {
	err := j.Cmd.Start()
	if err != nil {
		log.Panicf("cmd start error: %v\n", err)
	}
	log.Printf("job started\n")
	go timedSIGTERM(j.Cmd.Process, j.ExecPeriod)
	log.Printf("waiting for job terminate\n")
	err = j.Cmd.Wait()
	if err != nil {
		log.Panicf("cmd wait error: %v\n", err)
	}
	if j.Cmd.ProcessState.Success() {
		j.Res = DecodeStream(&j.OutBuf)
		//fmt.Printf("-->%v\n", j.Res)
		for _, e := range j.Res {
			tmp, _ := e.(map[string]interface{})
			s_pid, _ := tmp["pid"].(string)
			var e_pid int
			if len(s_pid) != 0 {
				e_pid, err = strconv.Atoi(s_pid)
				if err != nil {
					log.Fatalf("error convert atoi: %v\n", err)
				}
			}
			argList := GetArguments(e_pid)
			if len(argList) != 0 {
				if argList[0] == "/usr/libexec/qemu-kvm" {
					//fmt.Printf("%v %v %v\n", i, tmp, argList[2])
					ioread, _ := tmp["read"].(float64)
					iowrite, _ := tmp["write"].(float64)

					var storageType TiramisuStorage
					DB.First(&storageType, "vm_name = ?", argList[2])
					// fmt.Printf("\n\n---->[%v]\n\n", storageType)

					vminfo := RawVMData{
						VMName:    argList[2],
						IOPSRead:  ioread / j.ExecPeriod.Seconds(),
						IOPSWrite: iowrite / j.ExecPeriod.Seconds(),
						ISSSD:     storageType.CurrentPool == "SSD",
					}
					//fmt.Printf("\n\n---->[%v]\n\n", vminfo)
					//fmt.Printf("\n[%v]\n", vminfo)
					DB.Where(RawVMData{VMName: argList[2]}).Assign(vminfo).FirstOrCreate(&vminfo)
				}
			}
		}
		j.InitCmd()
		j.ExecTimed()
	}
}

func (j *JobScheduler) Execute() {
	if j.Type == Streaming {
		j.ExecStreaming()

	} else if j.Type == Timed {
		j.ExecTimed()
	}
}

// TODO: Refactor avg parameters as module.
func AssignAverage() {
	log.Printf("AssignAverage Invoked\n")
	var HDDVMList []RawVMData
	var SSDVMList []RawVMData
	DB.Model(&RawVMData{}).Where("isssd = ?", true).Find(&SSDVMList)
	DB.Model(&RawVMData{}).Where("isssd = ?", false).Find(&HDDVMList)

	var avgHDDLatency float64
	var avgSSDLatency float64
	var avgSSDIOPS float64
	var avgHDDIOPS float64

	for _, e := range HDDVMList {
		avgHDDIOPS += e.IOPSRead + e.IOPSWrite
		avgHDDLatency += e.LatencyRead + e.LatencyWrite
	}
	for _, e := range SSDVMList {
		avgSSDIOPS += e.IOPSRead + e.IOPSWrite
		avgSSDLatency += e.LatencyRead + e.LatencyWrite
	}
	avgHDDIOPS = (0.5 * avgHDDIOPS) / float64(len(HDDVMList))
	avgHDDLatency = (0.5 * avgHDDLatency) / float64(len(HDDVMList))
	avgSSDIOPS = (0.5 * avgSSDIOPS) / float64(len(SSDVMList))
	avgSSDLatency = (0.5 * avgSSDLatency) / float64(len(SSDVMList))

	for _, e := range HDDVMList {
		tiramisu_state := TiramisuState{
			IOPS:        (e.IOPSRead + e.IOPSWrite) / 2,
			Latency:     (e.LatencyRead + e.LatencyWrite) / 2,
			IOPS_HDD:    (e.IOPSRead + e.IOPSWrite) / 2,
			Latency_HDD: (e.LatencyRead + e.LatencyWrite) / 2,
			IOPS_SSD:    avgSSDIOPS,
			Latency_SSD: avgSSDLatency,
			Name:        e.VMName,
		}
		DB.Model(&TiramisuState{}).Where("vm_name = ?", e.VMName).Assign(tiramisu_state).FirstOrCreate(&tiramisu_state)
	}
	for _, e := range SSDVMList {
		tiramisu_state := TiramisuState{
			IOPS:        (e.IOPSRead + e.IOPSWrite) / 2,
			Latency:     (e.LatencyRead + e.LatencyWrite) / 2,
			IOPS_HDD:    avgHDDIOPS,
			Latency_HDD: avgHDDLatency,
			IOPS_SSD:    (e.IOPSRead + e.IOPSWrite) / 2,
			Latency_SSD: (e.LatencyRead + e.LatencyWrite) / 2,
			Name:        e.VMName,
		}
		DB.Model(&TiramisuState{}).Where("vm_name = ?", e.VMName).Assign(tiramisu_state).FirstOrCreate(&tiramisu_state)
	}
}
<|endoftext|>


<commit_before><commit_msg>working AES encryption/decryption<commit_after><|endoftext|>


<commit_before>package bolt

import (
	"bytes"
	"sort"
)

// Cursor represents an iterator that can traverse over all key/value pairs in a bucket in sorted order.
// Cursors see nested buckets with value == nil.
// Cursors can be obtained from a transaction and are valid as long as the transaction is open.
type Cursor struct {
	bucket *Bucket
	stack  []elemRef
}

// Bucket returns the bucket that this cursor was created from.
func (c *Cursor) Bucket() *Bucket {
	return c.bucket
}

// First moves the cursor to the first item in the bucket and returns its key and value.
// If the bucket is empty then a nil key and value are returned.
func (c *Cursor) First() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")
	c.stack = c.stack[:0]
	p, n := c.bucket.pageNode(c.bucket.root)
	c.stack = append(c.stack, elemRef{page: p, node: n, index: 0})
	c.first()
	k, v, flags := c.keyValue()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v

}

// Last moves the cursor to the last item in the bucket and returns its key and value.
// If the bucket is empty then a nil key and value are returned.
func (c *Cursor) Last() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")
	c.stack = c.stack[:0]
	p, n := c.bucket.pageNode(c.bucket.root)
	ref := elemRef{page: p, node: n}
	ref.index = ref.count() - 1
	c.stack = append(c.stack, ref)
	c.last()
	k, v, flags := c.keyValue()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Next moves the cursor to the next item in the bucket and returns its key and value.
// If the cursor is at the end of the bucket then a nil key and value are returned.
func (c *Cursor) Next() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")
	k, v, flags := c.next()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Prev moves the cursor to the previous item in the bucket and returns its key and value.
// If the cursor is at the beginning of the bucket then a nil key and value are returned.
func (c *Cursor) Prev() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")

	// Attempt to move back one element until we're successful.
	// Move up the stack as we hit the beginning of each page in our stack.
	for i := len(c.stack) - 1; i >= 0; i-- {
		elem := &c.stack[i]
		if elem.index > 0 {
			elem.index--
			break
		}
		c.stack = c.stack[:i]
	}

	// If we've hit the end then return nil.
	if len(c.stack) == 0 {
		return nil, nil
	}

	// Move down the stack to find the last element of the last leaf under this branch.
	c.last()
	k, v, flags := c.keyValue()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Seek moves the cursor to a given key and returns it.
// If the key does not exist then the next key is used. If no keys
// follow, a nil key is returned.
func (c *Cursor) Seek(seek []byte) (key []byte, value []byte) {
	k, v, flags := c.seek(seek)

	// If we ended up after the last element of a page then move to the next one.
	if ref := &c.stack[len(c.stack)-1]; ref.index >= ref.count() {
		k, v, flags = c.next()
	}

	if k == nil {
		return nil, nil
	} else if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Delete removes the current key/value under the cursor from the bucket.
// Delete fails if current key/value is a bucket or if the transaction is not writable.
func (c *Cursor) Delete() error {
	if c.bucket.tx.db == nil {
		return ErrTxClosed
	} else if !c.bucket.Writable() {
		return ErrTxNotWritable
	}

	key, _, flags := c.keyValue()
	// Return an error if current value is a bucket.
	if (flags & bucketLeafFlag) != 0 {
		return ErrIncompatibleValue
	}
	c.node().del(key)

	return nil
}

// seek moves the cursor to a given key and returns it.
// If the key does not exist then the next key is used.
func (c *Cursor) seek(seek []byte) (key []byte, value []byte, flags uint32) {
	_assert(c.bucket.tx.db != nil, "tx closed")

	// Start from root page/node and traverse to correct page.
	c.stack = c.stack[:0]
	c.search(seek, c.bucket.root)
	ref := &c.stack[len(c.stack)-1]

	// If the cursor is pointing to the end of page/node then return nil.
	if ref.index >= ref.count() {
		return nil, nil, 0
	}

	// If this is a bucket then return a nil value.
	return c.keyValue()
}

// first moves the cursor to the first leaf element under the last page in the stack.
func (c *Cursor) first() {
	for {
		// Exit when we hit a leaf page.
		var ref = &c.stack[len(c.stack)-1]
		if ref.isLeaf() {
			break
		}

		// Keep adding pages pointing to the first element to the stack.
		var pgid pgid
		if ref.node != nil {
			pgid = ref.node.inodes[ref.index].pgid
		} else {
			pgid = ref.page.branchPageElement(uint16(ref.index)).pgid
		}
		p, n := c.bucket.pageNode(pgid)
		c.stack = append(c.stack, elemRef{page: p, node: n, index: 0})
	}
}

// last moves the cursor to the last leaf element under the last page in the stack.
func (c *Cursor) last() {
	for {
		// Exit when we hit a leaf page.
		ref := &c.stack[len(c.stack)-1]
		if ref.isLeaf() {
			break
		}

		// Keep adding pages pointing to the last element in the stack.
		var pgid pgid
		if ref.node != nil {
			pgid = ref.node.inodes[ref.index].pgid
		} else {
			pgid = ref.page.branchPageElement(uint16(ref.index)).pgid
		}
		p, n := c.bucket.pageNode(pgid)

		var nextRef = elemRef{page: p, node: n}
		nextRef.index = nextRef.count() - 1
		c.stack = append(c.stack, nextRef)
	}
}

// next moves to the next leaf element and returns the key and value.
// If the cursor is at the last leaf element then it stays there and returns nil.
func (c *Cursor) next() (key []byte, value []byte, flags uint32) {
	// Attempt to move over one element until we're successful.
	// Move up the stack as we hit the end of each page in our stack.
	var i int
	for i = len(c.stack) - 1; i >= 0; i-- {
		elem := &c.stack[i]
		if elem.index < elem.count()-1 {
			elem.index++
			break
		}
	}

	// If we've hit the root page then stop and return. This will leave the
	// cursor on the last element of the last page.
	if i == -1 {
		return nil, nil, 0
	}

	// Otherwise start from where we left off in the stack and find the
	// first element of the first leaf page.
	c.stack = c.stack[:i+1]
	c.first()
	return c.keyValue()
}

// search recursively performs a binary search against a given page/node until it finds a given key.
func (c *Cursor) search(key []byte, pgid pgid) {
	p, n := c.bucket.pageNode(pgid)
	if p != nil {
		_assert((p.flags&(branchPageFlag|leafPageFlag)) != 0, "invalid page type: %d: %x", p.id, p.flags)
	}
	e := elemRef{page: p, node: n}
	c.stack = append(c.stack, e)

	// If we're on a leaf page/node then find the specific node.
	if e.isLeaf() {
		c.nsearch(key)
		return
	}

	if n != nil {
		c.searchNode(key, n)
		return
	}
	c.searchPage(key, p)
}

func (c *Cursor) searchNode(key []byte, n *node) {
	var exact bool
	index := sort.Search(len(n.inodes), func(i int) bool {
		// TODO(benbjohnson): Optimize this range search. It's a bit hacky right now.
		// sort.Search() finds the lowest index where f() != -1 but we need the highest index.
		ret := bytes.Compare(n.inodes[i].key, key)
		if ret == 0 {
			exact = true
		}
		return ret != -1
	})
	if !exact && index > 0 {
		index--
	}
	c.stack[len(c.stack)-1].index = index

	// Recursively search to the next page.
	c.search(key, n.inodes[index].pgid)
}

func (c *Cursor) searchPage(key []byte, p *page) {
	// Binary search for the correct range.
	inodes := p.branchPageElements()

	var exact bool
	index := sort.Search(int(p.count), func(i int) bool {
		// TODO(benbjohnson): Optimize this range search. It's a bit hacky right now.
		// sort.Search() finds the lowest index where f() != -1 but we need the highest index.
		ret := bytes.Compare(inodes[i].key(), key)
		if ret == 0 {
			exact = true
		}
		return ret != -1
	})
	if !exact && index > 0 {
		index--
	}
	c.stack[len(c.stack)-1].index = index

	// Recursively search to the next page.
	c.search(key, inodes[index].pgid)
}

// nsearch searches the leaf node on the top of the stack for a key.
func (c *Cursor) nsearch(key []byte) {
	e := &c.stack[len(c.stack)-1]
	p, n := e.page, e.node

	// If we have a node then search its inodes.
	if n != nil {
		index := sort.Search(len(n.inodes), func(i int) bool {
			return bytes.Compare(n.inodes[i].key, key) != -1
		})
		e.index = index
		return
	}

	// If we have a page then search its leaf elements.
	inodes := p.leafPageElements()
	index := sort.Search(int(p.count), func(i int) bool {
		return bytes.Compare(inodes[i].key(), key) != -1
	})
	e.index = index
}

// keyValue returns the key and value of the current leaf element.
func (c *Cursor) keyValue() ([]byte, []byte, uint32) {
	ref := &c.stack[len(c.stack)-1]
	if ref.count() == 0 || ref.index >= ref.count() {
		return nil, nil, 0
	}

	// Retrieve value from node.
	if ref.node != nil {
		inode := &ref.node.inodes[ref.index]
		return inode.key, inode.value, inode.flags
	}

	// Or retrieve value from page.
	elem := ref.page.leafPageElement(uint16(ref.index))
	return elem.key(), elem.value(), elem.flags
}

// node returns the node that the cursor is currently positioned on.
func (c *Cursor) node() *node {
	_assert(len(c.stack) > 0, "accessing a node with a zero-length cursor stack")

	// If the top of the stack is a leaf node then just return it.
	if ref := &c.stack[len(c.stack)-1]; ref.node != nil && ref.isLeaf() {
		return ref.node
	}

	// Start from root and traverse down the hierarchy.
	var n = c.stack[0].node
	if n == nil {
		n = c.bucket.node(c.stack[0].page.id, nil)
	}
	for _, ref := range c.stack[:len(c.stack)-1] {
		_assert(!n.isLeaf, "expected branch node")
		n = n.childAt(int(ref.index))
	}
	_assert(n.isLeaf, "expected leaf node")
	return n
}

// elemRef represents a reference to an element on a given page/node.
type elemRef struct {
	page  *page
	node  *node
	index int
}

// isLeaf returns whether the ref is pointing at a leaf page/node.
func (r *elemRef) isLeaf() bool {
	if r.node != nil {
		return r.node.isLeaf
	}
	return (r.page.flags & leafPageFlag) != 0
}

// count returns the number of inodes or page elements.
func (r *elemRef) count() int {
	if r.node != nil {
		return len(r.node.inodes)
	}
	return int(r.page.count)
}
<commit_msg>Clarify cursor behavior after mutation.<commit_after>package bolt

import (
	"bytes"
	"sort"
)

// Cursor represents an iterator that can traverse over all key/value pairs in a bucket in sorted order.
// Cursors see nested buckets with value == nil.
// Cursors can be obtained from a transaction and are valid as long as the transaction is open.
//
// Changing data while traversing with a cursor may cause it to be invalidated
// and return unexpected keys and/or values. You must reposition your cursor
// after mutating data.
type Cursor struct {
	bucket *Bucket
	stack  []elemRef
}

// Bucket returns the bucket that this cursor was created from.
func (c *Cursor) Bucket() *Bucket {
	return c.bucket
}

// First moves the cursor to the first item in the bucket and returns its key and value.
// If the bucket is empty then a nil key and value are returned.
func (c *Cursor) First() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")
	c.stack = c.stack[:0]
	p, n := c.bucket.pageNode(c.bucket.root)
	c.stack = append(c.stack, elemRef{page: p, node: n, index: 0})
	c.first()
	k, v, flags := c.keyValue()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v

}

// Last moves the cursor to the last item in the bucket and returns its key and value.
// If the bucket is empty then a nil key and value are returned.
func (c *Cursor) Last() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")
	c.stack = c.stack[:0]
	p, n := c.bucket.pageNode(c.bucket.root)
	ref := elemRef{page: p, node: n}
	ref.index = ref.count() - 1
	c.stack = append(c.stack, ref)
	c.last()
	k, v, flags := c.keyValue()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Next moves the cursor to the next item in the bucket and returns its key and value.
// If the cursor is at the end of the bucket then a nil key and value are returned.
func (c *Cursor) Next() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")
	k, v, flags := c.next()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Prev moves the cursor to the previous item in the bucket and returns its key and value.
// If the cursor is at the beginning of the bucket then a nil key and value are returned.
func (c *Cursor) Prev() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")

	// Attempt to move back one element until we're successful.
	// Move up the stack as we hit the beginning of each page in our stack.
	for i := len(c.stack) - 1; i >= 0; i-- {
		elem := &c.stack[i]
		if elem.index > 0 {
			elem.index--
			break
		}
		c.stack = c.stack[:i]
	}

	// If we've hit the end then return nil.
	if len(c.stack) == 0 {
		return nil, nil
	}

	// Move down the stack to find the last element of the last leaf under this branch.
	c.last()
	k, v, flags := c.keyValue()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Seek moves the cursor to a given key and returns it.
// If the key does not exist then the next key is used. If no keys
// follow, a nil key is returned.
func (c *Cursor) Seek(seek []byte) (key []byte, value []byte) {
	k, v, flags := c.seek(seek)

	// If we ended up after the last element of a page then move to the next one.
	if ref := &c.stack[len(c.stack)-1]; ref.index >= ref.count() {
		k, v, flags = c.next()
	}

	if k == nil {
		return nil, nil
	} else if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Delete removes the current key/value under the cursor from the bucket.
// Delete fails if current key/value is a bucket or if the transaction is not writable.
func (c *Cursor) Delete() error {
	if c.bucket.tx.db == nil {
		return ErrTxClosed
	} else if !c.bucket.Writable() {
		return ErrTxNotWritable
	}

	key, _, flags := c.keyValue()
	// Return an error if current value is a bucket.
	if (flags & bucketLeafFlag) != 0 {
		return ErrIncompatibleValue
	}
	c.node().del(key)

	return nil
}

// seek moves the cursor to a given key and returns it.
// If the key does not exist then the next key is used.
func (c *Cursor) seek(seek []byte) (key []byte, value []byte, flags uint32) {
	_assert(c.bucket.tx.db != nil, "tx closed")

	// Start from root page/node and traverse to correct page.
	c.stack = c.stack[:0]
	c.search(seek, c.bucket.root)
	ref := &c.stack[len(c.stack)-1]

	// If the cursor is pointing to the end of page/node then return nil.
	if ref.index >= ref.count() {
		return nil, nil, 0
	}

	// If this is a bucket then return a nil value.
	return c.keyValue()
}

// first moves the cursor to the first leaf element under the last page in the stack.
func (c *Cursor) first() {
	for {
		// Exit when we hit a leaf page.
		var ref = &c.stack[len(c.stack)-1]
		if ref.isLeaf() {
			break
		}

		// Keep adding pages pointing to the first element to the stack.
		var pgid pgid
		if ref.node != nil {
			pgid = ref.node.inodes[ref.index].pgid
		} else {
			pgid = ref.page.branchPageElement(uint16(ref.index)).pgid
		}
		p, n := c.bucket.pageNode(pgid)
		c.stack = append(c.stack, elemRef{page: p, node: n, index: 0})
	}
}

// last moves the cursor to the last leaf element under the last page in the stack.
func (c *Cursor) last() {
	for {
		// Exit when we hit a leaf page.
		ref := &c.stack[len(c.stack)-1]
		if ref.isLeaf() {
			break
		}

		// Keep adding pages pointing to the last element in the stack.
		var pgid pgid
		if ref.node != nil {
			pgid = ref.node.inodes[ref.index].pgid
		} else {
			pgid = ref.page.branchPageElement(uint16(ref.index)).pgid
		}
		p, n := c.bucket.pageNode(pgid)

		var nextRef = elemRef{page: p, node: n}
		nextRef.index = nextRef.count() - 1
		c.stack = append(c.stack, nextRef)
	}
}

// next moves to the next leaf element and returns the key and value.
// If the cursor is at the last leaf element then it stays there and returns nil.
func (c *Cursor) next() (key []byte, value []byte, flags uint32) {
	// Attempt to move over one element until we're successful.
	// Move up the stack as we hit the end of each page in our stack.
	var i int
	for i = len(c.stack) - 1; i >= 0; i-- {
		elem := &c.stack[i]
		if elem.index < elem.count()-1 {
			elem.index++
			break
		}
	}

	// If we've hit the root page then stop and return. This will leave the
	// cursor on the last element of the last page.
	if i == -1 {
		return nil, nil, 0
	}

	// Otherwise start from where we left off in the stack and find the
	// first element of the first leaf page.
	c.stack = c.stack[:i+1]
	c.first()
	return c.keyValue()
}

// search recursively performs a binary search against a given page/node until it finds a given key.
func (c *Cursor) search(key []byte, pgid pgid) {
	p, n := c.bucket.pageNode(pgid)
	if p != nil {
		_assert((p.flags&(branchPageFlag|leafPageFlag)) != 0, "invalid page type: %d: %x", p.id, p.flags)
	}
	e := elemRef{page: p, node: n}
	c.stack = append(c.stack, e)

	// If we're on a leaf page/node then find the specific node.
	if e.isLeaf() {
		c.nsearch(key)
		return
	}

	if n != nil {
		c.searchNode(key, n)
		return
	}
	c.searchPage(key, p)
}

func (c *Cursor) searchNode(key []byte, n *node) {
	var exact bool
	index := sort.Search(len(n.inodes), func(i int) bool {
		// TODO(benbjohnson): Optimize this range search. It's a bit hacky right now.
		// sort.Search() finds the lowest index where f() != -1 but we need the highest index.
		ret := bytes.Compare(n.inodes[i].key, key)
		if ret == 0 {
			exact = true
		}
		return ret != -1
	})
	if !exact && index > 0 {
		index--
	}
	c.stack[len(c.stack)-1].index = index

	// Recursively search to the next page.
	c.search(key, n.inodes[index].pgid)
}

func (c *Cursor) searchPage(key []byte, p *page) {
	// Binary search for the correct range.
	inodes := p.branchPageElements()

	var exact bool
	index := sort.Search(int(p.count), func(i int) bool {
		// TODO(benbjohnson): Optimize this range search. It's a bit hacky right now.
		// sort.Search() finds the lowest index where f() != -1 but we need the highest index.
		ret := bytes.Compare(inodes[i].key(), key)
		if ret == 0 {
			exact = true
		}
		return ret != -1
	})
	if !exact && index > 0 {
		index--
	}
	c.stack[len(c.stack)-1].index = index

	// Recursively search to the next page.
	c.search(key, inodes[index].pgid)
}

// nsearch searches the leaf node on the top of the stack for a key.
func (c *Cursor) nsearch(key []byte) {
	e := &c.stack[len(c.stack)-1]
	p, n := e.page, e.node

	// If we have a node then search its inodes.
	if n != nil {
		index := sort.Search(len(n.inodes), func(i int) bool {
			return bytes.Compare(n.inodes[i].key, key) != -1
		})
		e.index = index
		return
	}

	// If we have a page then search its leaf elements.
	inodes := p.leafPageElements()
	index := sort.Search(int(p.count), func(i int) bool {
		return bytes.Compare(inodes[i].key(), key) != -1
	})
	e.index = index
}

// keyValue returns the key and value of the current leaf element.
func (c *Cursor) keyValue() ([]byte, []byte, uint32) {
	ref := &c.stack[len(c.stack)-1]
	if ref.count() == 0 || ref.index >= ref.count() {
		return nil, nil, 0
	}

	// Retrieve value from node.
	if ref.node != nil {
		inode := &ref.node.inodes[ref.index]
		return inode.key, inode.value, inode.flags
	}

	// Or retrieve value from page.
	elem := ref.page.leafPageElement(uint16(ref.index))
	return elem.key(), elem.value(), elem.flags
}

// node returns the node that the cursor is currently positioned on.
func (c *Cursor) node() *node {
	_assert(len(c.stack) > 0, "accessing a node with a zero-length cursor stack")

	// If the top of the stack is a leaf node then just return it.
	if ref := &c.stack[len(c.stack)-1]; ref.node != nil && ref.isLeaf() {
		return ref.node
	}

	// Start from root and traverse down the hierarchy.
	var n = c.stack[0].node
	if n == nil {
		n = c.bucket.node(c.stack[0].page.id, nil)
	}
	for _, ref := range c.stack[:len(c.stack)-1] {
		_assert(!n.isLeaf, "expected branch node")
		n = n.childAt(int(ref.index))
	}
	_assert(n.isLeaf, "expected leaf node")
	return n
}

// elemRef represents a reference to an element on a given page/node.
type elemRef struct {
	page  *page
	node  *node
	index int
}

// isLeaf returns whether the ref is pointing at a leaf page/node.
func (r *elemRef) isLeaf() bool {
	if r.node != nil {
		return r.node.isLeaf
	}
	return (r.page.flags & leafPageFlag) != 0
}

// count returns the number of inodes or page elements.
func (r *elemRef) count() int {
	if r.node != nil {
		return len(r.node.inodes)
	}
	return int(r.page.count)
}
<|endoftext|>