<commit_before>package imageserver

import (
	"bytes"
	"encoding/binary"
	"fmt"
)

/*
Image represents a raw image.

Binary encoding

- Format length (uint32)

- Format (string)

- Data length (uint32)

- Data([]byte)

Numbers are encoded using little-endian order.
*/
type Image struct {
	Format string // png, jpeg, bmp, gif, ...
	Data   []byte // raw image data
}

// NewImageUnmarshalBinary creates a new Image from serialized bytes
func NewImageUnmarshalBinary(marshalledData []byte) (*Image, error) {
	image := new(Image)

	err := image.UnmarshalBinary(marshalledData)
	if err != nil {
		return nil, err
	}

	return image, nil
}

// MarshalBinary serializes the Image to bytes
//
// It's very unlikely that it returns an error (impossible?)
func (image *Image) MarshalBinary() ([]byte, error) {
	buffer := new(bytes.Buffer)

	formatLen := uint32(len(image.Format))
	binary.Write(buffer, binary.LittleEndian, &formatLen)
	buffer.Write([]byte(image.Format))

	dataLen := uint32(len(image.Data))
	binary.Write(buffer, binary.LittleEndian, &dataLen)
	buffer.Write(image.Data)

	return buffer.Bytes(), nil
}

// UnmarshalBinary unserializes bytes to the Image
func (image *Image) UnmarshalBinary(data []byte) error {
	dataPosition := 0
	readData := func(length int) ([]byte, error) {
		dataEnd := dataPosition + length
		if dataEnd > len(data) {
			return nil, &ImageError{Message: fmt.Sprintf("unmarshal: unexpected end of data at index %d instead of %d", len(data), dataEnd)}
		}
		d := data[dataPosition:dataEnd]
		dataPosition = dataEnd
		return d, nil
	}

	var imageFormatLength uint32
	if d, err := readData(4); err == nil {
		binary.Read(bytes.NewReader(d), binary.LittleEndian, &imageFormatLength)
	} else {
		return err
	}

	if d, err := readData(int(imageFormatLength)); err == nil {
		image.Format = string(d)
	} else {
		return err
	}

	var imageDataLength uint32
	if d, err := readData(4); err == nil {
		binary.Read(bytes.NewReader(d), binary.LittleEndian, &imageDataLength)
	} else {
		return err
	}

	if d, err := readData(int(imageDataLength)); err == nil {
		image.Data = d
	} else {
		return err
	}

	return nil
}

// ImageEqual compares two images and returns true if they are equal
func ImageEqual(image1, image2 *Image) bool {
	if image1 == image2 {
		return true
	}

	if image1 == nil || image2 == nil {
		return false
	}

	if image1.Format != image2.Format {
		return false
	}

	if !bytes.Equal(image1.Data, image2.Data) {
		return false
	}

	return true
}

// ImageError in an Image error
type ImageError struct {
	Message string
}

func (err *ImageError) Error() string {
	return fmt.Sprintf("image error: %s", err.Message)
}
<commit_msg>add Buffer Pool for Image marshaling<commit_after>package imageserver

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"sync"
)

/*
Image represents a raw image.

Binary encoding:

- Format length (uint32)

- Format (string)

- Data length (uint32)

- Data([]byte)

Numbers are encoded using little-endian order.
*/
type Image struct {
	Format string // png, jpeg, bmp, gif, ...
	Data   []byte // raw image data
}

// NewImageUnmarshalBinary creates a new Image from serialized bytes.
//
// The caller must not reuse data after that.
func NewImageUnmarshalBinary(marshalledData []byte) (*Image, error) {
	image := new(Image)
	err := image.UnmarshalBinary(marshalledData)
	if err != nil {
		return nil, err
	}
	return image, nil
}

var bufferPool = &sync.Pool{
	New: func() interface{} {
		return new(bytes.Buffer)
	},
}

// MarshalBinary serializes the Image to bytes.
//
// It's very unlikely that it returns an error. (impossible?)
func (image *Image) MarshalBinary() ([]byte, error) {
	buf := bufferPool.Get().(*bytes.Buffer)
	buf.Reset()

	formatLen := uint32(len(image.Format))
	binary.Write(buf, binary.LittleEndian, &formatLen)
	buf.Write([]byte(image.Format))

	dataLen := uint32(len(image.Data))
	binary.Write(buf, binary.LittleEndian, &dataLen)
	buf.Write(image.Data)

	b := buf.Bytes()
	bufferPool.Put(buf)
	return b, nil
}

// UnmarshalBinary unserializes bytes to the Image.
//
// The caller must not reuse data after that.
func (image *Image) UnmarshalBinary(data []byte) error {
	dataPosition := 0
	readData := func(length int) ([]byte, error) {
		dataEnd := dataPosition + length
		if dataEnd > len(data) {
			return nil, &ImageError{Message: fmt.Sprintf("unmarshal: unexpected end of data at index %d instead of %d", len(data), dataEnd)}
		}
		d := data[dataPosition:dataEnd]
		dataPosition = dataEnd
		return d, nil
	}

	var imageFormatLength uint32
	if d, err := readData(4); err == nil {
		binary.Read(bytes.NewReader(d), binary.LittleEndian, &imageFormatLength)
	} else {
		return err
	}

	if d, err := readData(int(imageFormatLength)); err == nil {
		image.Format = string(d)
	} else {
		return err
	}

	var imageDataLength uint32
	if d, err := readData(4); err == nil {
		binary.Read(bytes.NewReader(d), binary.LittleEndian, &imageDataLength)
	} else {
		return err
	}

	if d, err := readData(int(imageDataLength)); err == nil {
		image.Data = d
	} else {
		return err
	}

	return nil
}

// ImageEqual compares two images and returns true if they are equal.
func ImageEqual(image1, image2 *Image) bool {
	if image1 == image2 {
		return true
	}

	if image1 == nil || image2 == nil {
		return false
	}

	if image1.Format != image2.Format {
		return false
	}

	if !bytes.Equal(image1.Data, image2.Data) {
		return false
	}

	return true
}

// ImageError is an Image error.
type ImageError struct {
	Message string
}

func (err *ImageError) Error() string {
	return fmt.Sprintf("image error: %s", err.Message)
}
<|endoftext|>


<commit_before>/*
 * imgor
 * Copyright (c) 2012 Jimmy Zelinskie
 * Licensed under the MIT license.
 */

package main

import (
	"errors"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"os"
	"text/template"
)

// Globals
var (
	uploadTemplate *template.Template
	errorTemplate  *template.Template
	imgdir         string
)

// Check for errors
func check(err error) {
	if err != nil {
		panic(err)
	}
}

// MIME Validator
func validateimage(h *multipart.FileHeader) (ext string, err error) {
	mimeArray := h.Header["Content-Type"]
	mime := mimeArray[0]
	if mime == "image/jpeg" {
		ext = ".jpg"
	} else if mime == "image/png" {
		ext = ".png"
	} else {
		err = errors.New("Unsupported filetype uploaded")
	}
	return
}

// Upload Handler
func upload(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		// Load the upload page if they aren't posting an image
		uploadTemplate.Execute(w, nil)
		return
	}

	// Get image from POST
	f, h, err := r.FormFile("image")
	defer f.Close()
	check(err)

	// Check MIME and get a file extension
	_, err = validateimage(h)
	check(err)

	// Read and write the uploaded file to disk
	filename := imgdir + h.Filename
	filebytes, err := ioutil.ReadAll(f)
	check(err)
	err = ioutil.WriteFile(filename, filebytes, 0744)
	check(err)

	// Redirect to the view page
	http.Redirect(w, r, "/view?id="+filename[6:], http.StatusFound)
}

// View Handler
func view(w http.ResponseWriter, r *http.Request) {
	var filename string
	filename = imgdir + r.FormValue("id")

	if filename[len(filename)-4:] == ".jpg" {
		w.Header().Set("Content-Type", "image/jpeg")
	} else if filename[len(filename)-4:] == ".png" {
		w.Header().Set("Content-Type", "image/png")
	}

	http.ServeFile(w, r, filename)
}

// Error page's variables
type ErrorPage struct {
	Error error
}

// One clean error page
func errorHandler(fn http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if e, ok := recover().(error); ok {
				contents := ErrorPage{Error: e}
				w.WriteHeader(http.StatusInternalServerError)
				errorTemplate.Execute(w, contents)
			}
		}()
		fn(w, r)
	}
}

func main() {
	var err error

	// Set imgdir and make sure it exists!
	imgdir = "./img/"
	_ = os.Mkdir(imgdir[2 : len(imgdir)-1], 0744)

	// Load up templates and check for errors
	uploadTemplate, err = template.ParseFiles("upload.html")
	check(err)
	errorTemplate, err = template.ParseFiles("error.html")
	check(err)

	http.HandleFunc("/", errorHandler(upload))
	http.HandleFunc("/view", errorHandler(view))
	http.ListenAndServe(":3000", nil)
}
<commit_msg>sha1 random name, filetypes required<commit_after>/*
 * imgor
 * Copyright (c) 2012 Jimmy Zelinskie
 * Licensed under the MIT license.
 */

package main

import (
	"errors"
  "crypto/sha1"
	"io/ioutil"
	"mime/multipart"
  "fmt"
	"net/http"
	"os"
	"text/template"
)

// Globals
var (
	uploadTemplate *template.Template
	errorTemplate  *template.Template
	imgdir         string
)

// Check for errors
func check(err error) {
	if err != nil {
		panic(err)
	}
}

func generatefilename(d []byte) string {
  sha := sha1.New()
  return fmt.Sprintf("%x", string(sha.Sum(d))[0:10])
}

// MIME Validator
func validateimage(h *multipart.FileHeader) (ext string, err error) {
	mimeArray := h.Header["Content-Type"]
	mime := mimeArray[0]
	if mime == "image/jpeg" {
		ext = ".jpg"
	} else if mime == "image/png" {
		ext = ".png"
	} else {
		err = errors.New("Unsupported filetype uploaded")
	}
	return
}

// Upload Handler
func upload(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		// Load the upload page if they aren't posting an image
		uploadTemplate.Execute(w, nil)
		return
	}

	// Get image from POST
	f, h, err := r.FormFile("image")
	defer f.Close()
	check(err)

	// Check MIME and get a file extension
  ext, err := validateimage(h)
	check(err)

	// Read and write the uploaded file to disk
	filebytes, err := ioutil.ReadAll(f)
	check(err)
  filename := imgdir + generatefilename(filebytes) + ext
	err = ioutil.WriteFile(filename, filebytes, 0744)
	check(err)

	// Redirect to the view page
	http.Redirect(w, r, "/view?id="+filename[6:], http.StatusFound)
}

// View Handler
func view(w http.ResponseWriter, r *http.Request) {
	var filename string
	filename = imgdir + r.FormValue("id")

  if filename[len(filename)-3:] == "jpg" {
		w.Header().Set("Content-Type", "image/jpeg")
  } else if filename[len(filename)-3:] == "png" {
		w.Header().Set("Content-Type", "image/png")
	} else {
    panic(errors.New("No supported filetype specified"))
  }

	http.ServeFile(w, r, filename)
}

// Error page's variables
type ErrorPage struct {
	Error error
}

// One clean error page
func errorHandler(fn http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if e, ok := recover().(error); ok {
				contents := ErrorPage{Error: e}
				w.WriteHeader(http.StatusInternalServerError)
				errorTemplate.Execute(w, contents)
			}
		}()
		fn(w, r)
	}
}

func main() {
	var err error

	// Set imgdir and make sure it exists!
	imgdir = "./img/"
	_ = os.Mkdir(imgdir[2 : len(imgdir)-1], 0744)

	// Load up templates and check for errors
	uploadTemplate, err = template.ParseFiles("upload.html")
	check(err)
	errorTemplate, err = template.ParseFiles("error.html")
	check(err)

	http.HandleFunc("/", errorHandler(upload))
	http.HandleFunc("/view", errorHandler(view))
	http.ListenAndServe(":3000", nil)
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"database/sql"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"

	_ "github.com/mattn/go-sqlite3"
)

const (
	diskIndexDB = "./file-index.db"
	breakPoint  = 10000
)

var db *sql.DB

func main() {
	//init
	os.Remove(diskIndexDB)
	os.Create("indexer.txt")
	os.Create("copy.txt")

	bootDatabase()

	diskA := os.Args[1]
	diskB := os.Args[2]
	dest := os.Args[3]

	getFilesFromFolder(true, diskA)
	getFilesFromFolder(false, diskB)

	fmt.Println("--- FINISHED INDEXING ---")
	fmt.Println("--- STARTING COPY ---")

	files := findMissingFiles()
	copyFiles(dest, files)

	db.Close()

	fmt.Println("--- FINISHED ---")

}

func bootDatabase() {
	var err error

	//We're going to store disk A information in a SQLite database
	db, err = sql.Open("sqlite3", diskIndexDB)
	if err != nil {
		log.Fatal("Could not create the SQLite database", err)
	}

	_, err = db.Exec(`
	CREATE TABLE file (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		origin BOOLEAN,
		hash VARCHAR(64),
        path VARCHAR(2048)
    );`)

	if err != nil {
		log.Fatal("DB-TABLE CREATE: fail ", err)
	}
}

func insertFile(origin bool, filePath string) {
	//Not really sure what to do if there is an error, so let's log and quit to be safe. We don't want to mis files
	stmt, err := db.Prepare("INSERT INTO file(origin, hash, path) VALUES (?,?,?)")
	if err != nil {
		redoLog("indexer", filePath)
		return
	}

	md5, err := getMD5FromFile(filePath)
	if err != nil {
		redoLog("indexer", filePath)
		return
	}

	_, err = stmt.Exec(origin, md5, filePath)
	if err != nil {
		redoLog("indexer", filePath)
		return
	}
}

func getFilesFromFolder(origin bool, searchDir string) {
	fmt.Println("Indexing folder: " + searchDir)

	var count int64 = 0

	filepath.Walk(searchDir, func(path string, f os.FileInfo, err error) error {
		if !f.IsDir() {
			count++
			if count%breakPoint == 0 {
				fmt.Println(count)
			}
			insertFile(origin, path)
		}

		return nil
	})
}

func getMD5FromFile(filePath string) (string, error) {
	f, err := os.Open(filePath)
	if err != nil {
		return "", err
	}
	defer f.Close()

	h := md5.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}

	return fmt.Sprintf("%x", h.Sum(nil)), nil
}

func findMissingFiles() []string {
	//We left join to get the missing files and retrieve only the ones from the source
	missingQuery := `
	SELECT source.path,
		destination.id AS missing
	FROM   FILE source
		LEFT JOIN FILE destination
				ON source.hash = destination.hash
					AND source.origin = 1
					AND destination.origin = 0
	WHERE  missing IS NULL
		AND source.origin = 1`

	rows, err := db.Query(missingQuery)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	files := []string{}
	for rows.Next() {
		var path string
		var weDoNotCare interface{}
		err = rows.Scan(&path, &weDoNotCare)

		if err != nil {
			log.Fatal(err)
		}

		files = append(files, path)
	}

	err = rows.Err()
	if err != nil {
		log.Fatal(err)
	}

	return files

}

func copyFiles(dir string, files []string) {
	err := os.MkdirAll(dir, os.ModePerm)
	if err != nil {
		log.Fatal("COPY INIT: Could not create copy-directory")
	}

	lengthFiles := len(files)
	count := 0
	for _, file := range files {
		count++
		percentage := int((count / lengthFiles) * 100)

		if percentage%2 == 0 {
			fmt.Println(fmt.Sprintf("%d%%", percentage))
		}

		err = copyFile(file, dir+file)

		if err != nil {
			redoLog("copy", file)
			fmt.Println(err)
		}
	}
}

func copyFile(src string, dst string) (err error) {
	//Create output dir
	dir := filepath.Dir(dst)
	os.MkdirAll(dir, os.ModePerm)

	//Do the copy
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.Create(dst)
	if err != nil {
		return err
	}

	defer func() {
		cerr := out.Close()
		if err == nil {
			err = cerr
		}
	}()

	if _, err = io.Copy(out, in); err != nil {
		return err
	}

	return out.Sync()
}

//Dirty function to write the filepaths that crashedd to a file for after-processing
func redoLog(tag string, path string) {
	f, _ := os.OpenFile(tag+".txt", os.O_APPEND|os.O_WRONLY, os.ModeAppend)
	f.WriteString(path + "\r\n") //Window line endings, since we're running on that os for the backup
	f.Close()
}
<commit_msg>More windows hacks<commit_after>package main

import (
	"crypto/md5"
	"database/sql"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"

	_ "github.com/mattn/go-sqlite3"
)

const (
	diskIndexDB = "./file-index.db"
	breakPoint  = 10000
)

var db *sql.DB

func main() {
	//init
	os.Remove(diskIndexDB)
	os.Create("indexer.txt")
	os.Create("copy.txt")

	bootDatabase()

	diskA := os.Args[1]
	diskB := os.Args[2]
	dest := os.Args[3]

	getFilesFromFolder(true, diskA)
	getFilesFromFolder(false, diskB)

	fmt.Println("--- FINISHED INDEXING ---")
	fmt.Println("--- STARTING COPY ---")

	files := findMissingFiles()
	copyFiles(dest, files)

	db.Close()

	fmt.Println("--- FINISHED ---")

}

func bootDatabase() {
	var err error

	//We're going to store disk A information in a SQLite database
	db, err = sql.Open("sqlite3", diskIndexDB)
	if err != nil {
		log.Fatal("Could not create the SQLite database", err)
	}

	_, err = db.Exec(`
	CREATE TABLE file (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		origin BOOLEAN,
		hash VARCHAR(64),
        path VARCHAR(2048)
    );`)

	if err != nil {
		log.Fatal("DB-TABLE CREATE: fail ", err)
	}
}

func insertFile(origin bool, filePath string) {
	//Not really sure what to do if there is an error, so let's log and quit to be safe. We don't want to mis files
	stmt, err := db.Prepare("INSERT INTO file(origin, hash, path) VALUES (?,?,?)")
	if err != nil {
		redoLog("indexer", filePath)
		return
	}

	md5, err := getMD5FromFile(filePath)
	if err != nil {
		redoLog("indexer", filePath)
		return
	}

	_, err = stmt.Exec(origin, md5, filePath)
	if err != nil {
		redoLog("indexer", filePath)
		return
	}
}

func getFilesFromFolder(origin bool, searchDir string) {
	fmt.Println("Indexing folder: " + searchDir)

	var count int64 = 0

	filepath.Walk(searchDir, func(path string, f os.FileInfo, err error) error {
		if !f.IsDir() {
			count++
			if count%breakPoint == 0 {
				fmt.Println(count)
			}
			insertFile(origin, path)
		}

		return nil
	})
}

func getMD5FromFile(filePath string) (string, error) {
	f, err := os.Open(filePath)
	if err != nil {
		return "", err
	}
	defer f.Close()

	h := md5.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}

	return fmt.Sprintf("%x", h.Sum(nil)), nil
}

func findMissingFiles() []string {
	//We left join to get the missing files and retrieve only the ones from the source
	missingQuery := `
	SELECT source.path,
		destination.id AS missing
	FROM   FILE source
		LEFT JOIN FILE destination
				ON source.hash = destination.hash
					AND source.origin = 1
					AND destination.origin = 0
	WHERE  missing IS NULL
		AND source.origin = 1`

	rows, err := db.Query(missingQuery)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	files := []string{}
	for rows.Next() {
		var path string
		var weDoNotCare interface{}
		err = rows.Scan(&path, &weDoNotCare)

		if err != nil {
			log.Fatal(err)
		}

		files = append(files, path)
	}

	err = rows.Err()
	if err != nil {
		log.Fatal(err)
	}

	return files

}

func copyFiles(dir string, files []string) {
	err := os.MkdirAll(dir, os.ModePerm)
	if err != nil {
		log.Fatal("COPY INIT: Could not create copy-directory")
	}

	lengthFiles := len(files)
	count := 0
	for _, file := range files {
		count++
		percentage := int((count / lengthFiles) * 100)

		if percentage%2 == 0 {
			fmt.Println(fmt.Sprintf("%d%%", percentage))
		}

		err = copyFile(file, strings.Replace(dir+file, ":\\", "", -1)) //Windows does not like ex. C:\ in the path as new folder, so hack it away

		if err != nil {
			redoLog("copy", file)
		}
	}
}

func copyFile(src string, dst string) (err error) {
	//Create output dir
	dir := filepath.Dir(dst)
	os.MkdirAll(dir, os.ModePerm)

	//Do the copy
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.Create(dst)
	if err != nil {
		return err
	}

	defer func() {
		cerr := out.Close()
		if err == nil {
			err = cerr
		}
	}()

	if _, err = io.Copy(out, in); err != nil {
		return err
	}

	return out.Sync()
}

//Dirty function to write the filepaths that crashedd to a file for after-processing
func redoLog(tag string, path string) {
	f, _ := os.OpenFile(tag+".txt", os.O_APPEND|os.O_WRONLY, os.ModeAppend)
	f.WriteString(path + "\r\n") //Window line endings, since we're running on that os for the backup
	f.Close()
}
<|endoftext|>


<commit_before>package job

import (
	"errors"
	"github.com/watermint/toolbox/domain/model/mo_path"
	"github.com/watermint/toolbox/infra/control/app_control"
	"github.com/watermint/toolbox/infra/control/app_workflow"
	"github.com/watermint/toolbox/infra/ui/app_msg"
	"github.com/watermint/toolbox/infra/util/ut_process"
	"github.com/watermint/toolbox/quality/infra/qt_endtoend"
	"go.uber.org/zap"
	"os"
	"os/exec"
	"time"
)

type Run struct {
	Fork                    bool
	TimeoutSeconds          int
	RunbookPath             mo_path.FileSystemPath
	ErrorRunBookNotFound    app_msg.Message
	ErrorTimeoutRequireFork app_msg.Message
	ErrorUnableToFork       app_msg.Message
}

func (z *Run) execFork(c app_control.Control) error {
	ui := c.UI()
	l := c.Log()

	l.Info("Fork", zap.String("cmd", os.Args[0]), zap.String("runbook", z.RunbookPath.Path()))
	cmd := exec.Command(os.Args[0], "job", "run", z.RunbookPath.Path())
	pl := ut_process.NewLogger(cmd, c)
	pl.Start()
	defer pl.Close()

	l.Debug("Start")
	err := cmd.Start()
	if err != nil {
		ui.Error(z.ErrorUnableToFork.With("Error", err))
		return err
	}

	if z.TimeoutSeconds < 1 {
		l.Info("Waiting for finish process")
		if err := cmd.Wait(); err != nil {
			l.Info("The process finished with an error", zap.Error(err))
		} else {
			l.Info("The process finished")
		}
		return err
	}

	running := true
	go func() {
		cmd.Wait()
		l.Debug("The process finished")
		running = false
	}()

	timeout := time.Now().Add(time.Duration(z.TimeoutSeconds) * time.Second)
	l.Info("Waiting for process", zap.String("timeout", timeout.Format(time.RFC3339)))
	for {
		time.Sleep(500 * time.Microsecond)
		if !running {
			return nil
		}
		if time.Now().After(timeout) {
			l.Debug("Execution timeout, try send kill signal to the process")
			err = cmd.Process.Kill()
			l.Debug("Signal sent", zap.Error(err))
			cmd.Process.Release()
			return nil
		}
	}
}

func (z *Run) execInProcess(c app_control.Control) error {
	ui := c.UI()
	rb, found := app_workflow.NewRunBook(z.RunbookPath.Path())
	if !found {
		ui.Error(z.ErrorRunBookNotFound.With("Path", z.RunbookPath.Path()))
		return errors.New("runbook not found")
	}
	if err := rb.Verify(c); err != nil {
		c.Log().Debug("Verification failure")
		return err
	}
	return rb.Run(c)
}

func (z *Run) Exec(c app_control.Control) error {
	ui := c.UI()
	if !z.Fork && z.TimeoutSeconds > 0 {
		ui.Error(z.ErrorTimeoutRequireFork)
		return errors.New("-timeout-seconds option requires fork")
	}

	if z.Fork {
		return z.execFork(c)
	}
	return z.execInProcess(c)
}

func (z *Run) Test(c app_control.Control) error {
	// Can't test from this func. Test on tbx_test
	return qt_endtoend.ScenarioTest()
}

func (z *Run) Preset() {
	z.Fork = false
	z.TimeoutSeconds = 0
}
<commit_msg>fix #279 : fix for fork<commit_after>package job

import (
	"errors"
	"github.com/watermint/toolbox/domain/model/mo_path"
	"github.com/watermint/toolbox/infra/control/app_control"
	"github.com/watermint/toolbox/infra/control/app_workflow"
	"github.com/watermint/toolbox/infra/ui/app_msg"
	"github.com/watermint/toolbox/infra/util/ut_process"
	"github.com/watermint/toolbox/quality/infra/qt_endtoend"
	"go.uber.org/zap"
	"os"
	"os/exec"
	"time"
)

type Run struct {
	Fork                    bool
	TimeoutSeconds          int
	RunbookPath             mo_path.FileSystemPath
	ErrorRunBookNotFound    app_msg.Message
	ErrorTimeoutRequireFork app_msg.Message
	ErrorUnableToFork       app_msg.Message
}

func (z *Run) execFork(c app_control.Control) error {
	ui := c.UI()
	l := c.Log()

	l.Info("Fork", zap.String("cmd", os.Args[0]), zap.String("runbook", z.RunbookPath.Path()))
	cmd := exec.Command(os.Args[0], "job", "run", "-runbook-path", z.RunbookPath.Path())
	pl := ut_process.NewLogger(cmd, c)
	pl.Start()
	defer pl.Close()

	l.Debug("Start")
	err := cmd.Start()
	if err != nil {
		ui.Error(z.ErrorUnableToFork.With("Error", err))
		return err
	}

	if z.TimeoutSeconds < 1 {
		l.Info("Waiting for finish process")
		if err := cmd.Wait(); err != nil {
			l.Info("The process finished with an error", zap.Error(err))
		} else {
			l.Info("The process finished")
		}
		return err
	}

	running := true
	go func() {
		cmd.Wait()
		l.Debug("The process finished")
		running = false
	}()

	timeout := time.Now().Add(time.Duration(z.TimeoutSeconds) * time.Second)
	l.Info("Waiting for process", zap.String("timeout", timeout.Format(time.RFC3339)))
	for {
		time.Sleep(500 * time.Microsecond)
		if !running {
			return nil
		}
		if time.Now().After(timeout) {
			l.Debug("Execution timeout, try send kill signal to the process")
			err = cmd.Process.Kill()
			l.Debug("Signal sent", zap.Error(err))
			cmd.Process.Release()
			return nil
		}
	}
}

func (z *Run) execInProcess(c app_control.Control) error {
	ui := c.UI()
	rb, found := app_workflow.NewRunBook(z.RunbookPath.Path())
	if !found {
		ui.Error(z.ErrorRunBookNotFound.With("Path", z.RunbookPath.Path()))
		return errors.New("runbook not found")
	}
	if err := rb.Verify(c); err != nil {
		c.Log().Debug("Verification failure")
		return err
	}
	return rb.Run(c)
}

func (z *Run) Exec(c app_control.Control) error {
	ui := c.UI()
	if !z.Fork && z.TimeoutSeconds > 0 {
		ui.Error(z.ErrorTimeoutRequireFork)
		return errors.New("-timeout-seconds option requires fork")
	}

	if z.Fork {
		return z.execFork(c)
	}
	return z.execInProcess(c)
}

func (z *Run) Test(c app_control.Control) error {
	// Can't test from this func. Test on tbx_test
	return qt_endtoend.ScenarioTest()
}

func (z *Run) Preset() {
	z.Fork = false
	z.TimeoutSeconds = 0
}
<|endoftext|>


<commit_before>package router

import (
	"encoding/json"
	"errors"
	"html"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/NyaaPantsu/nyaa/config"
	"github.com/NyaaPantsu/nyaa/db"
	"github.com/NyaaPantsu/nyaa/model"
	"github.com/NyaaPantsu/nyaa/service"
	"github.com/NyaaPantsu/nyaa/service/api"
	"github.com/NyaaPantsu/nyaa/service/torrent"
	"github.com/NyaaPantsu/nyaa/service/upload"
	"github.com/NyaaPantsu/nyaa/service/user"
	"github.com/NyaaPantsu/nyaa/util/log"
	"github.com/NyaaPantsu/nyaa/util/search"
	"github.com/gorilla/mux"
)

// APIHandler : Controller for api request on torrent list
func APIHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	page := vars["page"]
	whereParams := serviceBase.WhereParams{}
	req := apiService.TorrentsRequest{}
	defer r.Body.Close()

	contentType := r.Header.Get("Content-Type")
	if contentType == "application/json" {
		d := json.NewDecoder(r.Body)
		if err := d.Decode(&req); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
		}

		if req.MaxPerPage == 0 {
			req.MaxPerPage = config.Conf.Navigation.TorrentsPerPage
		}
		if req.Page <= 0 {
			req.Page = 1
		}

		whereParams = req.ToParams()
	} else {
		var err error
		maxString := r.URL.Query().Get("max")
		if maxString != "" {
			req.MaxPerPage, err = strconv.Atoi(maxString)
			if !log.CheckError(err) {
				req.MaxPerPage = config.Conf.Navigation.TorrentsPerPage
			}
		} else {
			req.MaxPerPage = config.Conf.Navigation.TorrentsPerPage
		}

		req.Page = 1
		if page != "" {
			req.Page, err = strconv.Atoi(html.EscapeString(page))
			if !log.CheckError(err) {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			if req.Page <= 0 {
				NotFoundHandler(w, r)
				return
			}
		}
	}

	torrents, nbTorrents, err := torrentService.GetTorrents(whereParams, req.MaxPerPage, req.MaxPerPage*(req.Page-1))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	b := model.APIResultJSON{
		Torrents: model.TorrentsToJSON(torrents),
	}
	b.QueryRecordCount = req.MaxPerPage
	b.TotalRecordCount = nbTorrents
	w.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w).Encode(b)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

// APIViewHandler : Controller for viewing a torrent by its ID
func APIViewHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]

	torrent, err := torrentService.GetTorrentByID(id)

	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}

	b := torrent.ToJSON()
	w.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w).Encode(b)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer r.Body.Close()
}

// APIViewHeadHandler : Controller for checking a torrent by its ID
func APIViewHeadHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id, err := strconv.ParseInt(vars["id"], 10, 32)
	defer r.Body.Close()
	if err != nil {
		return
	}

	_, err = torrentService.GetRawTorrentByID(uint(id))

	if err != nil {
		NotFoundHandler(w, r)
		return
	}

	w.Write(nil)
}

// APIUploadHandler : Controller for uploading a torrent with api
func APIUploadHandler(w http.ResponseWriter, r *http.Request) {
	token := r.Header.Get("Authorization")
	user, _, _, err := userService.RetrieveUserByAPIToken(token)
	defer r.Body.Close()

	if err != nil {
		http.Error(w, "Error API token doesn't exist", http.StatusBadRequest)
		return
	}

	if !uploadService.IsUploadEnabled(user) {
		http.Error(w, "Error uploads are disabled", http.StatusBadRequest)
		return
	}

	if user.ID == 0 {
		http.Error(w, apiService.ErrAPIKey.Error(), http.StatusUnauthorized)
		return
	}

	upload := apiService.TorrentRequest{}
	contentType := r.Header.Get("Content-Type")
	if contentType != "application/json" && !strings.HasPrefix(contentType, "multipart/form-data") && r.Header.Get("Content-Type") != "application/x-www-form-urlencoded" {
		// TODO What should we do here ? upload is empty so we shouldn't
		// create a torrent from it
		http.Error(w, errors.New("Please provide either of Content-Type: application/json header or multipart/form-data").Error(), http.StatusInternalServerError)
		return
	}
	// As long as the right content-type is sent, formValue is smart enough to parse it
	err = upload.ExtractInfo(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	status := model.TorrentStatusNormal
	if upload.Remake { // overrides trusted
		status = model.TorrentStatusRemake
	} else if user.IsTrusted() {
		status = model.TorrentStatusTrusted
	}

	err = torrentService.ExistOrDelete(upload.Infohash, user)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	torrent := model.Torrent{
		Name:        upload.Name,
		Category:    upload.CategoryID,
		SubCategory: upload.SubCategoryID,
		Status:      status,
		Hidden:      upload.Hidden,
		Hash:        upload.Infohash,
		Date:        time.Now(),
		Filesize:    upload.Filesize,
		Description: upload.Description,
		WebsiteLink: upload.WebsiteLink,
		UploaderID:  user.ID}
	torrent.ParseTrackers(upload.Trackers)
	db.ORM.Create(&torrent)

	if db.ElasticSearchClient != nil {
		err := torrent.AddToESIndex(db.ElasticSearchClient)
		if err == nil {
			log.Infof("Successfully added torrent to ES index.")
		} else {
			log.Errorf("Unable to add torrent to ES index: %s", err)
		}
	} else {
		log.Errorf("Unable to create elasticsearch client: %s", err)
	}

	torrentService.NewTorrentEvent(Router, user, &torrent)
	// add filelist to files db, if we have one
	if len(upload.FileList) > 0 {
		for _, uploadedFile := range upload.FileList {
			file := model.File{TorrentID: torrent.ID, Filesize: upload.Filesize}
			err := file.SetPath(uploadedFile.Path)
			if err != nil {
				http.Error(w, err.Error(), http.StatusBadRequest)
				return
			}
			db.ORM.Create(&file)
		}
	}
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
}

// APIUpdateHandler : Controller for updating a torrent with api
// FIXME Impossible to update a torrent uploaded by user 0
func APIUpdateHandler(w http.ResponseWriter, r *http.Request) {
	token := r.Header.Get("Authorization")
	user, _, _, err := userService.RetrieveUserByAPIToken(token)
	defer r.Body.Close()

	if err != nil {
		http.Error(w, "Error API token doesn't exist", http.StatusBadRequest)
		return
	}

	if !uploadService.IsUploadEnabled(user) {
		http.Error(w, "Error uploads are disabled", http.StatusBadRequest)
		return
	}

	if user.ID == 0 {
		http.Error(w, apiService.ErrAPIKey.Error(), http.StatusForbidden)
		return
	}
	contentType := r.Header.Get("Content-Type")
	if contentType != "application/json" && !strings.HasPrefix(contentType, "multipart/form-data") && r.Header.Get("Content-Type") != "application/x-www-form-urlencoded" {
		// TODO What should we do here ? upload is empty so we shouldn't
		// create a torrent from it
		http.Error(w, errors.New("Please provide either of Content-Type: application/json header or multipart/form-data").Error(), http.StatusInternalServerError)
		return
	}
	update := apiService.UpdateRequest{}
	err = update.Update.ExtractEditInfo(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	torrent := model.Torrent{}
	db.ORM.Where("torrent_id = ?", r.FormValue("id")).First(&torrent)
	if torrent.ID == 0 {
		http.Error(w, apiService.ErrTorrentID.Error(), http.StatusBadRequest)
		return
	}
	if torrent.UploaderID != 0 && torrent.UploaderID != user.ID { //&& user.Status != mod
		http.Error(w, apiService.ErrRights.Error(), http.StatusForbidden)
		return
	}
	update.UpdateTorrent(&torrent, user)
	torrentService.UpdateTorrent(torrent)
}

// APISearchHandler : Controller for searching with api
func APISearchHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	page := vars["page"]
	defer r.Body.Close()

	// db params url
	var err error
	pagenum := 1
	if page != "" {
		pagenum, err = strconv.Atoi(html.EscapeString(page))
		if !log.CheckError(err) {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if pagenum <= 0 {
			NotFoundHandler(w, r)
			return
		}
	}

	_, torrents, _, err := search.SearchByQuery(r, pagenum)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	b := model.TorrentsToJSON(torrents)
	w.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w).Encode(b)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}
<commit_msg>Fix the fact that we didn't have uploader name on api search<commit_after>package router

import (
	"encoding/json"
	"errors"
	"html"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/NyaaPantsu/nyaa/config"
	"github.com/NyaaPantsu/nyaa/db"
	"github.com/NyaaPantsu/nyaa/model"
	"github.com/NyaaPantsu/nyaa/service"
	"github.com/NyaaPantsu/nyaa/service/api"
	"github.com/NyaaPantsu/nyaa/service/torrent"
	"github.com/NyaaPantsu/nyaa/service/upload"
	"github.com/NyaaPantsu/nyaa/service/user"
	"github.com/NyaaPantsu/nyaa/util/log"
	"github.com/NyaaPantsu/nyaa/util/search"
	"github.com/gorilla/mux"
)

// APIHandler : Controller for api request on torrent list
func APIHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	page := vars["page"]
	whereParams := serviceBase.WhereParams{}
	req := apiService.TorrentsRequest{}
	defer r.Body.Close()

	contentType := r.Header.Get("Content-Type")
	if contentType == "application/json" {
		d := json.NewDecoder(r.Body)
		if err := d.Decode(&req); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
		}

		if req.MaxPerPage == 0 {
			req.MaxPerPage = config.Conf.Navigation.TorrentsPerPage
		}
		if req.Page <= 0 {
			req.Page = 1
		}

		whereParams = req.ToParams()
	} else {
		var err error
		maxString := r.URL.Query().Get("max")
		if maxString != "" {
			req.MaxPerPage, err = strconv.Atoi(maxString)
			if !log.CheckError(err) {
				req.MaxPerPage = config.Conf.Navigation.TorrentsPerPage
			}
		} else {
			req.MaxPerPage = config.Conf.Navigation.TorrentsPerPage
		}

		req.Page = 1
		if page != "" {
			req.Page, err = strconv.Atoi(html.EscapeString(page))
			if !log.CheckError(err) {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			if req.Page <= 0 {
				NotFoundHandler(w, r)
				return
			}
		}
	}

	torrents, nbTorrents, err := torrentService.GetTorrents(whereParams, req.MaxPerPage, req.MaxPerPage*(req.Page-1))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	b := model.APIResultJSON{
		Torrents: model.TorrentsToJSON(torrents),
	}
	b.QueryRecordCount = req.MaxPerPage
	b.TotalRecordCount = nbTorrents
	w.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w).Encode(b)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

// APIViewHandler : Controller for viewing a torrent by its ID
func APIViewHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]

	torrent, err := torrentService.GetTorrentByID(id)

	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}

	b := torrent.ToJSON()
	w.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w).Encode(b)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer r.Body.Close()
}

// APIViewHeadHandler : Controller for checking a torrent by its ID
func APIViewHeadHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id, err := strconv.ParseInt(vars["id"], 10, 32)
	defer r.Body.Close()
	if err != nil {
		return
	}

	_, err = torrentService.GetRawTorrentByID(uint(id))

	if err != nil {
		NotFoundHandler(w, r)
		return
	}

	w.Write(nil)
}

// APIUploadHandler : Controller for uploading a torrent with api
func APIUploadHandler(w http.ResponseWriter, r *http.Request) {
	token := r.Header.Get("Authorization")
	user, _, _, err := userService.RetrieveUserByAPIToken(token)
	defer r.Body.Close()

	if err != nil {
		http.Error(w, "Error API token doesn't exist", http.StatusBadRequest)
		return
	}

	if !uploadService.IsUploadEnabled(user) {
		http.Error(w, "Error uploads are disabled", http.StatusBadRequest)
		return
	}

	if user.ID == 0 {
		http.Error(w, apiService.ErrAPIKey.Error(), http.StatusUnauthorized)
		return
	}

	upload := apiService.TorrentRequest{}
	contentType := r.Header.Get("Content-Type")
	if contentType != "application/json" && !strings.HasPrefix(contentType, "multipart/form-data") && r.Header.Get("Content-Type") != "application/x-www-form-urlencoded" {
		// TODO What should we do here ? upload is empty so we shouldn't
		// create a torrent from it
		http.Error(w, errors.New("Please provide either of Content-Type: application/json header or multipart/form-data").Error(), http.StatusInternalServerError)
		return
	}
	// As long as the right content-type is sent, formValue is smart enough to parse it
	err = upload.ExtractInfo(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	status := model.TorrentStatusNormal
	if upload.Remake { // overrides trusted
		status = model.TorrentStatusRemake
	} else if user.IsTrusted() {
		status = model.TorrentStatusTrusted
	}

	err = torrentService.ExistOrDelete(upload.Infohash, user)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	torrent := model.Torrent{
		Name:        upload.Name,
		Category:    upload.CategoryID,
		SubCategory: upload.SubCategoryID,
		Status:      status,
		Hidden:      upload.Hidden,
		Hash:        upload.Infohash,
		Date:        time.Now(),
		Filesize:    upload.Filesize,
		Description: upload.Description,
		WebsiteLink: upload.WebsiteLink,
		UploaderID:  user.ID}
	torrent.ParseTrackers(upload.Trackers)
	db.ORM.Create(&torrent)

	if db.ElasticSearchClient != nil {
		err := torrent.AddToESIndex(db.ElasticSearchClient)
		if err == nil {
			log.Infof("Successfully added torrent to ES index.")
		} else {
			log.Errorf("Unable to add torrent to ES index: %s", err)
		}
	} else {
		log.Errorf("Unable to create elasticsearch client: %s", err)
	}

	torrentService.NewTorrentEvent(Router, user, &torrent)
	// add filelist to files db, if we have one
	if len(upload.FileList) > 0 {
		for _, uploadedFile := range upload.FileList {
			file := model.File{TorrentID: torrent.ID, Filesize: upload.Filesize}
			err := file.SetPath(uploadedFile.Path)
			if err != nil {
				http.Error(w, err.Error(), http.StatusBadRequest)
				return
			}
			db.ORM.Create(&file)
		}
	}
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
}

// APIUpdateHandler : Controller for updating a torrent with api
// FIXME Impossible to update a torrent uploaded by user 0
func APIUpdateHandler(w http.ResponseWriter, r *http.Request) {
	token := r.Header.Get("Authorization")
	user, _, _, err := userService.RetrieveUserByAPIToken(token)
	defer r.Body.Close()

	if err != nil {
		http.Error(w, "Error API token doesn't exist", http.StatusBadRequest)
		return
	}

	if !uploadService.IsUploadEnabled(user) {
		http.Error(w, "Error uploads are disabled", http.StatusBadRequest)
		return
	}

	if user.ID == 0 {
		http.Error(w, apiService.ErrAPIKey.Error(), http.StatusForbidden)
		return
	}
	contentType := r.Header.Get("Content-Type")
	if contentType != "application/json" && !strings.HasPrefix(contentType, "multipart/form-data") && r.Header.Get("Content-Type") != "application/x-www-form-urlencoded" {
		// TODO What should we do here ? upload is empty so we shouldn't
		// create a torrent from it
		http.Error(w, errors.New("Please provide either of Content-Type: application/json header or multipart/form-data").Error(), http.StatusInternalServerError)
		return
	}
	update := apiService.UpdateRequest{}
	err = update.Update.ExtractEditInfo(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	torrent := model.Torrent{}
	db.ORM.Where("torrent_id = ?", r.FormValue("id")).First(&torrent)
	if torrent.ID == 0 {
		http.Error(w, apiService.ErrTorrentID.Error(), http.StatusBadRequest)
		return
	}
	if torrent.UploaderID != 0 && torrent.UploaderID != user.ID { //&& user.Status != mod
		http.Error(w, apiService.ErrRights.Error(), http.StatusForbidden)
		return
	}
	update.UpdateTorrent(&torrent, user)
	torrentService.UpdateTorrent(torrent)
}

// APISearchHandler : Controller for searching with api
func APISearchHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	page := vars["page"]
	defer r.Body.Close()

	// db params url
	var err error
	pagenum := 1
	if page != "" {
		pagenum, err = strconv.Atoi(html.EscapeString(page))
		if !log.CheckError(err) {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if pagenum <= 0 {
			NotFoundHandler(w, r)
			return
		}
	}

	_, torrents, _, err := search.SearchByQueryWithUser(r, pagenum)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	b := model.TorrentsToJSON(torrents)
	w.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w).Encode(b)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/user"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

// Attrs store user-defined parameters
type Attrs struct {
	Region      string
	Bucket      string
	Config      string
	Section     string
	Concurrency int
}

// Get user-defined parameters from CLI. Also, this function provides
// some verbose output, for example, if bucket was not specified
func getArgs() *Attrs {
	var region, config string
	regionPtr := flag.String("region", "", "Defines region")
	bucketPtr := flag.String("bucket", "", "Defines bucket. default = empty")
	configPtr := flag.String("config", "", "Allow changing AWS account")
	sectionPtr := flag.String("section", "default", "Which part of AWS credentials to use")
	concurrencyPtr := flag.Int("maxcon", 10, "Set up maximum concurrency for this task. Default is 10")
	flag.Parse()
	if *bucketPtr == "" {
		fmt.Println("You haven't define bucket! Please, do it with -bucket= ")
		os.Exit(1)
	}
	if *regionPtr == "" {
		region = "us-east-1"
		fmt.Println("You haven't specified region. Default region will be us-east-1")
	} else {
		region = *regionPtr
	}
	if *configPtr == "" {
		usr, err := user.Current()
		if err != nil {
			panic(err)
		}
		config = usr.HomeDir + "/.aws/credentials"
	} else {
		config = *configPtr
	}
	attrs := Attrs{
		Region:      region,
		Bucket:      *bucketPtr,
		Config:      config,
		Section:     *sectionPtr,
		Concurrency: *concurrencyPtr,
	}
	return &attrs
}

func main() {
	attrs := getArgs()
	creds := credentials.NewSharedCredentials(attrs.Config, attrs.Section)
	_, err := creds.Get()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	// Create new connection to S3
	svc := s3.New(session.New(), &aws.Config{
		Region:      aws.String(attrs.Region),
		Credentials: creds,
	})
	params := &s3.ListObjectsInput{
		Bucket: aws.String(attrs.Bucket),
	}
	resp, _ := svc.ListObjects(params)
	fmt.Print(len(resp.Contents), " objects in the bucket.\n Processing... It could take a while...")

	// This is used to limit simultaneous goroutines
	throttle := make(chan int, attrs.Concurrency)
	var wg sync.WaitGroup

	// Loop trough the objects in the bucket and create a copy
	// of each object with the REDUCED_REDUNDANCY storage class
	for _, key := range resp.Contents {
		if *key.StorageClass != "REDUCED_REDUNDANCY" {
			throttle <- 1
			wg.Add(1)
			go func() {
				defer wg.Done()
				copyParams := &s3.CopyObjectInput{
					Bucket:       aws.String(attrs.Bucket),
					CopySource:   aws.String(attrs.Bucket + "/" + *key.Key),
					Key:          aws.String(*key.Key),
					StorageClass: aws.String("REDUCED_REDUNDANCY"),
				}
				_, err := svc.CopyObject(copyParams)
				if err != nil {
					panic(err)
				}
				fmt.Print(".")
				<-throttle
			}()
			wg.Wait()
		}
	}

	// Fill the channel to be sure, that all goroutines finished
	for i := 0; i < cap(throttle); i++ {
		throttle <- 1
	}
	fmt.Println("\nConversion done!")
}
<commit_msg>refactoring v1<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os/user"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"gopkg.in/cheggaaa/pb.v1"
)

// Attrs store user-defined parameters
type Attrs struct {
	Region      string
	Bucket      string
	Config      string
	Section     string
	Concurrency int
}

// Get user-defined parameters from CLI
var (
	regionPtr      = flag.String("region", "", "Defines region")
	bucketPtr      = flag.String("bucket", "", "Defines bucket. default = empty")
	configPtr      = flag.String("config", "", "Allow changing AWS account")
	sectionPtr     = flag.String("section", "default", "Which part of AWS credentials to use")
	concurrencyPtr = flag.Int("maxcon", 10, "Set up maximum concurrency for this task. Default is 10")
)

func convert(attrs Attrs) {
	creds := credentials.NewSharedCredentials(attrs.Config, attrs.Section)
	_, err := creds.Get()
	if err != nil {
		log.Fatal(err)
		return
	}
	// Create new connection to S3
	svc := s3.New(session.New(), &aws.Config{
		Region:      aws.String(attrs.Region),
		Credentials: creds,
	})
	params := &s3.ListObjectsInput{
		Bucket: aws.String(attrs.Bucket),
	}
	resp, _ := svc.ListObjects(params)
	fmt.Print(len(resp.Contents), " objects in the bucket.\n Processing... It could take a while...")

	// This is used to limit simultaneous goroutines
	throttle := make(chan int, attrs.Concurrency)
	var wg sync.WaitGroup

	// Loop trough the objects in the bucket and create a copy
	// of each object with the REDUCED_REDUNDANCY storage class
	bar := pb.StartNew(len(resp.Contents))
	for _, key := range resp.Contents {
		if *key.StorageClass != "REDUCED_REDUNDANCY" {
			throttle <- 1
			wg.Add(1)
			go func() {
				defer wg.Done()
				copyParams := &s3.CopyObjectInput{
					Bucket:       aws.String(attrs.Bucket),
					CopySource:   aws.String(attrs.Bucket + "/" + *key.Key),
					Key:          aws.String(*key.Key),
					StorageClass: aws.String("REDUCED_REDUNDANCY"),
				}
				_, err := svc.CopyObject(copyParams)
				if err != nil {
					log.Fatal(err)
				}

				<-throttle
			}()
			wg.Wait()
		}
		bar.Increment()
	}
	bar.FinishPrint("Conversion done!")
	// Fill the channel to be sure, that all goroutines finished
	for i := 0; i < cap(throttle); i++ {
		throttle <- 1
	}
}

func main() {
	var region, config string
	// Parsing arguments
	flag.Parse()
	if *bucketPtr == "" {
		log.Fatal("You haven't define bucket! Please, do it with -bucket= ")
		return
	}
	if *regionPtr == "" {
		region = "us-east-1"
		fmt.Println("You haven't specified region. Default region will be us-east-1")
	} else {
		region = *regionPtr
	}
	if *configPtr == "" {
		usr, err := user.Current()
		if err != nil {
			log.Fatal(err)
			return
		}
		config = usr.HomeDir + "/.aws/credentials"
	} else {
		config = *configPtr
	}
	attrs := Attrs{
		Region:      region,
		Bucket:      *bucketPtr,
		Config:      config,
		Section:     *sectionPtr,
		Concurrency: *concurrencyPtr,
	}

	convert(attrs)
}
<|endoftext|>


<commit_before>package ruler

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"

	"github.com/prometheus/common/log"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/rules"
)

// TODO: Extract configs client logic into go client library (ala users)

type configID int

type cortexConfig struct {
	RulesFiles map[string]string `json:"rules_files"`
}

type cortexConfigView struct {
	ConfigID configID     `json:"id"`
	Config   cortexConfig `json:"config"`
}

// cortexConfigsResponse is a response from server for getOrgConfigs
type cortexConfigsResponse struct {
	// Configs maps organization ID to their latest cortexConfigView.
	Configs map[string]cortexConfigView `json:"configs"`
}

func configsFromJSON(body io.Reader) (map[string]cortexConfigView, error) {
	var configs cortexConfigsResponse
	if err := json.NewDecoder(body).Decode(&configs); err != nil {
		log.Errorf("configs: couldn't decode JSON body: %v", err)
		return nil, err
	}
	log.Debugf("configs: got response: %v", configs)
	return configs.Configs, nil
}

// getLatestConfigID returns the last config ID from a set of configs.
func getLatestConfigID(configs map[string]cortexConfigView) configID {
	latest := configID(0)
	for _, config := range configs {
		if config.ConfigID > latest {
			latest = config.ConfigID
		}
	}
	return latest
}

// Get the rules from the cortex configuration.
//
// Strongly inspired by `loadGroups` in Prometheus.
func (c cortexConfig) GetRules() ([]rules.Rule, error) {
	result := []rules.Rule{}
	for fn, content := range c.RulesFiles {
		stmts, err := promql.ParseStmts(content)
		if err != nil {
			return nil, fmt.Errorf("error parsing %s: %s", fn, err)
		}

		for _, stmt := range stmts {
			var rule rules.Rule

			switch r := stmt.(type) {
			case *promql.AlertStmt:
				rule = rules.NewAlertingRule(r.Name, r.Expr, r.Duration, r.Labels, r.Annotations)

			case *promql.RecordStmt:
				rule = rules.NewRecordingRule(r.Name, r.Expr, r.Labels)

			default:
				return nil, fmt.Errorf("ruler.GetRules: unknown statement type")
			}
			result = append(result, rule)
		}
	}
	return result, nil
}

type configsAPI struct {
	url *url.URL
}

// getOrgConfigs returns all Cortex configurations from a configs api server
// that have been updated since the given time.
func (c *configsAPI) getOrgConfigs(since configID) (map[string]cortexConfigView, error) {
	suffix := ""
	if since != 0 {
		suffix = fmt.Sprintf("?since=%d", since)
	}
	url := fmt.Sprintf("%s/private/api/configs/org/cortex%s", c.url.String(), suffix)
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	client := &http.Client{}
	res, err := client.Do(req)
	defer res.Body.Close()
	if err != nil {
		return nil, err
	}
	if res.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("Invalid response from configs server: %v", res.StatusCode)
	}
	return configsFromJSON(res.Body)
}
<commit_msg>Move body close<commit_after>package ruler

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"

	"github.com/prometheus/common/log"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/rules"
)

// TODO: Extract configs client logic into go client library (ala users)

type configID int

type cortexConfig struct {
	RulesFiles map[string]string `json:"rules_files"`
}

type cortexConfigView struct {
	ConfigID configID     `json:"id"`
	Config   cortexConfig `json:"config"`
}

// cortexConfigsResponse is a response from server for getOrgConfigs
type cortexConfigsResponse struct {
	// Configs maps organization ID to their latest cortexConfigView.
	Configs map[string]cortexConfigView `json:"configs"`
}

func configsFromJSON(body io.Reader) (map[string]cortexConfigView, error) {
	var configs cortexConfigsResponse
	if err := json.NewDecoder(body).Decode(&configs); err != nil {
		log.Errorf("configs: couldn't decode JSON body: %v", err)
		return nil, err
	}
	log.Debugf("configs: got response: %v", configs)
	return configs.Configs, nil
}

// getLatestConfigID returns the last config ID from a set of configs.
func getLatestConfigID(configs map[string]cortexConfigView) configID {
	latest := configID(0)
	for _, config := range configs {
		if config.ConfigID > latest {
			latest = config.ConfigID
		}
	}
	return latest
}

// Get the rules from the cortex configuration.
//
// Strongly inspired by `loadGroups` in Prometheus.
func (c cortexConfig) GetRules() ([]rules.Rule, error) {
	result := []rules.Rule{}
	for fn, content := range c.RulesFiles {
		stmts, err := promql.ParseStmts(content)
		if err != nil {
			return nil, fmt.Errorf("error parsing %s: %s", fn, err)
		}

		for _, stmt := range stmts {
			var rule rules.Rule

			switch r := stmt.(type) {
			case *promql.AlertStmt:
				rule = rules.NewAlertingRule(r.Name, r.Expr, r.Duration, r.Labels, r.Annotations)

			case *promql.RecordStmt:
				rule = rules.NewRecordingRule(r.Name, r.Expr, r.Labels)

			default:
				return nil, fmt.Errorf("ruler.GetRules: unknown statement type")
			}
			result = append(result, rule)
		}
	}
	return result, nil
}

type configsAPI struct {
	url *url.URL
}

// getOrgConfigs returns all Cortex configurations from a configs api server
// that have been updated since the given time.
func (c *configsAPI) getOrgConfigs(since configID) (map[string]cortexConfigView, error) {
	suffix := ""
	if since != 0 {
		suffix = fmt.Sprintf("?since=%d", since)
	}
	url := fmt.Sprintf("%s/private/api/configs/org/cortex%s", c.url.String(), suffix)
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	client := &http.Client{}
	res, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("Invalid response from configs server: %v", res.StatusCode)
	}
	return configsFromJSON(res.Body)
}
<|endoftext|>


<commit_before>package ast

import (
	"errors"
	"github.com/matttproud/prometheus/model"
	"log"
	"math"
	"strings"
	"time"
)

// ----------------------------------------------------------------------------
// Raw data value types.

type Vector []*model.Sample
type Matrix []*model.SampleSet

type groupedAggregation struct {
	labels     model.Metric
	value      model.SampleValue
	groupCount int
}

type labelValuePair struct {
	label model.LabelName
	value model.LabelValue
}

// ----------------------------------------------------------------------------
// Enums.

// Rule language expression types.
type ExprType int

const (
	SCALAR ExprType = iota
	VECTOR
	MATRIX
	STRING
)

// Binary operator types.
type BinOpType int

const (
	ADD BinOpType = iota
	SUB
	MUL
	DIV
	MOD
	NE
	EQ
	GT
	LT
	GE
	LE
	AND
	OR
)

// Aggregation types.
type AggrType int

const (
	SUM AggrType = iota
	AVG
	MIN
	MAX
)

// ----------------------------------------------------------------------------
// Interfaces.

// All node interfaces include the Node interface.
type Node interface {
	Type() ExprType
	NodeTreeToDotGraph() string
}

// All node types implement one of the following interfaces. The name of the
// interface represents the type returned to the parent node.
type ScalarNode interface {
	Node
	Eval(timestamp *time.Time) model.SampleValue
}

type VectorNode interface {
	Node
	Eval(timestamp *time.Time) Vector
}

type MatrixNode interface {
	Node
	Eval(timestamp *time.Time) Matrix
	EvalBoundaries(timestamp *time.Time) Matrix
}

type StringNode interface {
	Node
	Eval(timestamp *time.Time) string
}

// ----------------------------------------------------------------------------
// ScalarNode types.

type (
	// A numeric literal.
	ScalarLiteral struct {
		value model.SampleValue
	}

	// A function of numeric return type.
	ScalarFunctionCall struct {
		function *Function
		args     []Node
	}

	// An arithmetic expression of numeric type.
	ScalarArithExpr struct {
		opType BinOpType
		lhs    ScalarNode
		rhs    ScalarNode
	}
)

// ----------------------------------------------------------------------------
// VectorNode types.

type (
	// Vector literal, i.e. metric name plus labelset.
	VectorLiteral struct {
		labels model.LabelSet
	}

	// A function of vector return type.
	VectorFunctionCall struct {
		function *Function
		args     []Node
	}

	// A vector aggregation with vector return type.
	VectorAggregation struct {
		aggrType AggrType
		groupBy  []model.LabelName
		vector   VectorNode
	}

	// An arithmetic expression of vector type.
	VectorArithExpr struct {
		opType BinOpType
		lhs    VectorNode
		rhs    Node
	}
)

// ----------------------------------------------------------------------------
// MatrixNode types.

type (
	// Matrix literal, i.e. metric name plus labelset and timerange.
	MatrixLiteral struct {
		labels   model.LabelSet
		interval time.Duration
	}
)

// ----------------------------------------------------------------------------
// StringNode types.

type (
	// String literal.
	StringLiteral struct {
		str string
	}

	// A function of string return type.
	StringFunctionCall struct {
		function *Function
		args     []Node
	}
)

// ----------------------------------------------------------------------------
// Implementations.

func (node ScalarLiteral) Type() ExprType      { return SCALAR }
func (node ScalarFunctionCall) Type() ExprType { return SCALAR }
func (node ScalarArithExpr) Type() ExprType    { return SCALAR }
func (node VectorLiteral) Type() ExprType      { return VECTOR }
func (node VectorFunctionCall) Type() ExprType { return VECTOR }
func (node VectorAggregation) Type() ExprType  { return VECTOR }
func (node VectorArithExpr) Type() ExprType    { return VECTOR }
func (node MatrixLiteral) Type() ExprType      { return MATRIX }
func (node StringLiteral) Type() ExprType      { return STRING }
func (node StringFunctionCall) Type() ExprType { return STRING }

func (node *ScalarLiteral) Eval(timestamp *time.Time) model.SampleValue {
	return node.value
}

func (node *ScalarArithExpr) Eval(timestamp *time.Time) model.SampleValue {
	lhs := node.lhs.Eval(timestamp)
	rhs := node.rhs.Eval(timestamp)
	return evalScalarBinop(node.opType, lhs, rhs)
}

func (node *ScalarFunctionCall) Eval(timestamp *time.Time) model.SampleValue {
	return node.function.callFn(timestamp, node.args).(model.SampleValue)
}

func (node *VectorAggregation) labelsToGroupingKey(labels model.Metric) string {
	keyParts := []string{}
	for _, keyLabel := range node.groupBy {
		keyParts = append(keyParts, string(labels[keyLabel]))
	}
	return strings.Join(keyParts, ",") // TODO not safe when label value contains comma.
}

func labelIntersection(metric1, metric2 model.Metric) model.Metric {
	intersection := model.Metric{}
	for label, value := range metric1 {
		if metric2[label] == value {
			intersection[label] = value
		}
	}
	return intersection
}

func (node *VectorAggregation) groupedAggregationsToVector(aggregations map[string]*groupedAggregation, timestamp *time.Time) Vector {
	vector := Vector{}
	for _, aggregation := range aggregations {
		if node.aggrType == AVG {
			aggregation.value = aggregation.value / model.SampleValue(aggregation.groupCount)
		}
		sample := &model.Sample{
			Metric:    aggregation.labels,
			Value:     aggregation.value,
			Timestamp: *timestamp,
		}
		vector = append(vector, sample)
	}
	return vector
}

func (node *VectorAggregation) Eval(timestamp *time.Time) Vector {
	vector := node.vector.Eval(timestamp)
	result := map[string]*groupedAggregation{}
	for _, sample := range vector {
		groupingKey := node.labelsToGroupingKey(sample.Metric)
		if groupedResult, ok := result[groupingKey]; ok {
			groupedResult.labels = labelIntersection(groupedResult.labels, sample.Metric)
			switch node.aggrType {
			case SUM:
				groupedResult.value += sample.Value
			case AVG:
				groupedResult.value += sample.Value
				groupedResult.groupCount++
			case MAX:
				if groupedResult.value < sample.Value {
					groupedResult.value = sample.Value
				}
			case MIN:
				if groupedResult.value > sample.Value {
					groupedResult.value = sample.Value
				}
			}
		} else {
			result[groupingKey] = &groupedAggregation{
				labels:     sample.Metric,
				value:      sample.Value,
				groupCount: 1,
			}
		}
	}
	return node.groupedAggregationsToVector(result, timestamp)
}

func (node *VectorLiteral) Eval(timestamp *time.Time) Vector {
	values, err := persistence.GetValueAtTime(node.labels, timestamp, &stalenessPolicy)
	if err != nil {
		log.Printf("Unable to get vector values")
		return Vector{}
	}
	return values
}

func (node *VectorFunctionCall) Eval(timestamp *time.Time) Vector {
	return node.function.callFn(timestamp, node.args).(Vector)
}

func evalScalarBinop(opType BinOpType,
	lhs model.SampleValue,
	rhs model.SampleValue) model.SampleValue {
	switch opType {
	case ADD:
		return lhs + rhs
	case SUB:
		return lhs - rhs
	case MUL:
		return lhs * rhs
	case DIV:
		if rhs != 0 {
			return lhs / rhs
		} else {
			return model.SampleValue(math.Inf(int(rhs)))
		}
	case MOD:
		if rhs != 0 {
			return model.SampleValue(int(lhs) % int(rhs))
		} else {
			return model.SampleValue(math.Inf(int(rhs)))
		}
	case EQ:
		if lhs == rhs {
			return 1
		} else {
			return 0
		}
	case NE:
		if lhs != rhs {
			return 1
		} else {
			return 0
		}
	case GT:
		if lhs > rhs {
			return 1
		} else {
			return 0
		}
	case LT:
		if lhs < rhs {
			return 1
		} else {
			return 0
		}
	case GE:
		if lhs >= rhs {
			return 1
		} else {
			return 0
		}
	case LE:
		if lhs <= rhs {
			return 1
		} else {
			return 0
		}
	}
	panic("Not all enum values enumerated in switch")
}

func evalVectorBinop(opType BinOpType,
	lhs model.SampleValue,
	rhs model.SampleValue) (model.SampleValue, bool) {
	switch opType {
	case ADD:
		return lhs + rhs, true
	case SUB:
		return lhs - rhs, true
	case MUL:
		return lhs * rhs, true
	case DIV:
		if rhs != 0 {
			return lhs / rhs, true
		} else {
			return model.SampleValue(math.Inf(int(rhs))), true
		}
	case MOD:
		if rhs != 0 {
			return model.SampleValue(int(lhs) % int(rhs)), true
		} else {
			return model.SampleValue(math.Inf(int(rhs))), true
		}
	case EQ:
		if lhs == rhs {
			return lhs, true
		} else {
			return 0, false
		}
	case NE:
		if lhs != rhs {
			return lhs, true
		} else {
			return 0, false
		}
	case GT:
		if lhs > rhs {
			return lhs, true
		} else {
			return 0, false
		}
	case LT:
		if lhs < rhs {
			return lhs, true
		} else {
			return 0, false
		}
	case GE:
		if lhs >= rhs {
			return lhs, true
		} else {
			return 0, false
		}
	case LE:
		if lhs <= rhs {
			return lhs, true
		} else {
			return 0, false
		}
	case AND:
		return lhs, true
	}
	panic("Not all enum values enumerated in switch")
}

func labelsEqual(labels1, labels2 model.Metric) bool {
	if len(labels1) != len(labels2) {
		return false
	}
	for label, value := range labels1 {
		if labels2[label] != value && label != "name" {
			return false
		}
	}
	return true
}

func (node *VectorArithExpr) Eval(timestamp *time.Time) Vector {
	lhs := node.lhs.Eval(timestamp)
	result := Vector{}
	if node.rhs.Type() == SCALAR {
		rhs := node.rhs.(ScalarNode).Eval(timestamp)
		for _, lhsSample := range lhs {
			value, keep := evalVectorBinop(node.opType, lhsSample.Value, rhs)
			if keep {
				lhsSample.Value = value
				result = append(result, lhsSample)
			}
		}
		return result
	} else if node.rhs.Type() == VECTOR {
		rhs := node.rhs.(VectorNode).Eval(timestamp)
		for _, lhsSample := range lhs {
			for _, rhsSample := range rhs {
				if labelsEqual(lhsSample.Metric, rhsSample.Metric) {
					value, keep := evalVectorBinop(node.opType, lhsSample.Value, rhsSample.Value)
					if keep {
						lhsSample.Value = value
						result = append(result, lhsSample)
					}
				}
			}
		}
		return result
	}
	panic("Invalid vector arithmetic expression operands")
}

func (node *MatrixLiteral) Eval(timestamp *time.Time) Matrix {
	values, err := persistence.GetRangeValues(node.labels, &model.Interval{}, &stalenessPolicy)
	if err != nil {
		log.Printf("Unable to get values for vector interval")
		return Matrix{}
	}
	return values
}

func (node *MatrixLiteral) EvalBoundaries(timestamp *time.Time) Matrix {
	interval := &model.Interval{
		OldestInclusive: timestamp.Add(-node.interval),
		NewestInclusive: *timestamp,
	}
	values, err := persistence.GetBoundaryValues(node.labels, interval, &stalenessPolicy)
	if err != nil {
		log.Printf("Unable to get boundary values for vector interval")
		return Matrix{}
	}
	return values
}

func (node *StringLiteral) Eval(timestamp *time.Time) string {
	return node.str
}

func (node *StringFunctionCall) Eval(timestamp *time.Time) string {
	return node.function.callFn(timestamp, node.args).(string)
}

// ----------------------------------------------------------------------------
// Constructors.

func NewScalarLiteral(value model.SampleValue) *ScalarLiteral {
	return &ScalarLiteral{
		value: value,
	}
}

func NewVectorLiteral(labels model.LabelSet) *VectorLiteral {
	return &VectorLiteral{
		labels: labels,
	}
}

func NewVectorAggregation(aggrType AggrType, vector VectorNode, groupBy []model.LabelName) *VectorAggregation {
	return &VectorAggregation{
		aggrType: aggrType,
		groupBy:  groupBy,
		vector:   vector,
	}
}

func NewFunctionCall(function *Function, args []Node) (Node, error) {
	if err := function.CheckArgTypes(args); err != nil {
		return nil, err
	}
	switch function.returnType {
	case SCALAR:
		return &ScalarFunctionCall{
			function: function,
			args:     args,
		}, nil
	case VECTOR:
		return &VectorFunctionCall{
			function: function,
			args:     args,
		}, nil
	case STRING:
		return &StringFunctionCall{
			function: function,
			args:     args,
		}, nil
	}
	panic("Function with invalid return type")
}

func nodesHaveTypes(nodes []Node, exprTypes []ExprType) bool {
	for _, node := range nodes {
		for _, exprType := range exprTypes {
			if node.Type() == exprType {
				return true
			}
		}
	}
	return false
}

func NewArithExpr(opType BinOpType, lhs Node, rhs Node) (Node, error) {
	if !nodesHaveTypes([]Node{lhs, rhs}, []ExprType{SCALAR, VECTOR}) {
		return nil, errors.New("Binary operands must be of vector or scalar type")
	}
	if lhs.Type() == SCALAR && rhs.Type() == VECTOR {
		return nil, errors.New("Left side of vector binary operation must be of vector type")
	}

	if opType == AND || opType == OR {
		if lhs.Type() == SCALAR || rhs.Type() == SCALAR {
			return nil, errors.New("AND and OR operators may only be used between vectors")
		}
	}

	if lhs.Type() == VECTOR || rhs.Type() == VECTOR {
		return &VectorArithExpr{
			opType: opType,
			lhs:    lhs.(VectorNode),
			rhs:    rhs,
		}, nil
	}

	return &ScalarArithExpr{
		opType: opType,
		lhs:    lhs.(ScalarNode),
		rhs:    rhs.(ScalarNode),
	}, nil
}

func NewMatrixLiteral(vector *VectorLiteral, interval time.Duration) *MatrixLiteral {
	return &MatrixLiteral{
		labels:   vector.labels,
		interval: interval,
	}
}

func NewStringLiteral(str string) *StringLiteral {
	return &StringLiteral{
		str: str,
	}
}
<commit_msg>Set correct interval in MatrixLiteral.Eval().<commit_after>package ast

import (
	"errors"
	"github.com/matttproud/prometheus/model"
	"log"
	"math"
	"strings"
	"time"
)

// ----------------------------------------------------------------------------
// Raw data value types.

type Vector []*model.Sample
type Matrix []*model.SampleSet

type groupedAggregation struct {
	labels     model.Metric
	value      model.SampleValue
	groupCount int
}

type labelValuePair struct {
	label model.LabelName
	value model.LabelValue
}

// ----------------------------------------------------------------------------
// Enums.

// Rule language expression types.
type ExprType int

const (
	SCALAR ExprType = iota
	VECTOR
	MATRIX
	STRING
)

// Binary operator types.
type BinOpType int

const (
	ADD BinOpType = iota
	SUB
	MUL
	DIV
	MOD
	NE
	EQ
	GT
	LT
	GE
	LE
	AND
	OR
)

// Aggregation types.
type AggrType int

const (
	SUM AggrType = iota
	AVG
	MIN
	MAX
)

// ----------------------------------------------------------------------------
// Interfaces.

// All node interfaces include the Node interface.
type Node interface {
	Type() ExprType
	NodeTreeToDotGraph() string
}

// All node types implement one of the following interfaces. The name of the
// interface represents the type returned to the parent node.
type ScalarNode interface {
	Node
	Eval(timestamp *time.Time) model.SampleValue
}

type VectorNode interface {
	Node
	Eval(timestamp *time.Time) Vector
}

type MatrixNode interface {
	Node
	Eval(timestamp *time.Time) Matrix
	EvalBoundaries(timestamp *time.Time) Matrix
}

type StringNode interface {
	Node
	Eval(timestamp *time.Time) string
}

// ----------------------------------------------------------------------------
// ScalarNode types.

type (
	// A numeric literal.
	ScalarLiteral struct {
		value model.SampleValue
	}

	// A function of numeric return type.
	ScalarFunctionCall struct {
		function *Function
		args     []Node
	}

	// An arithmetic expression of numeric type.
	ScalarArithExpr struct {
		opType BinOpType
		lhs    ScalarNode
		rhs    ScalarNode
	}
)

// ----------------------------------------------------------------------------
// VectorNode types.

type (
	// Vector literal, i.e. metric name plus labelset.
	VectorLiteral struct {
		labels model.LabelSet
	}

	// A function of vector return type.
	VectorFunctionCall struct {
		function *Function
		args     []Node
	}

	// A vector aggregation with vector return type.
	VectorAggregation struct {
		aggrType AggrType
		groupBy  []model.LabelName
		vector   VectorNode
	}

	// An arithmetic expression of vector type.
	VectorArithExpr struct {
		opType BinOpType
		lhs    VectorNode
		rhs    Node
	}
)

// ----------------------------------------------------------------------------
// MatrixNode types.

type (
	// Matrix literal, i.e. metric name plus labelset and timerange.
	MatrixLiteral struct {
		labels   model.LabelSet
		interval time.Duration
	}
)

// ----------------------------------------------------------------------------
// StringNode types.

type (
	// String literal.
	StringLiteral struct {
		str string
	}

	// A function of string return type.
	StringFunctionCall struct {
		function *Function
		args     []Node
	}
)

// ----------------------------------------------------------------------------
// Implementations.

func (node ScalarLiteral) Type() ExprType      { return SCALAR }
func (node ScalarFunctionCall) Type() ExprType { return SCALAR }
func (node ScalarArithExpr) Type() ExprType    { return SCALAR }
func (node VectorLiteral) Type() ExprType      { return VECTOR }
func (node VectorFunctionCall) Type() ExprType { return VECTOR }
func (node VectorAggregation) Type() ExprType  { return VECTOR }
func (node VectorArithExpr) Type() ExprType    { return VECTOR }
func (node MatrixLiteral) Type() ExprType      { return MATRIX }
func (node StringLiteral) Type() ExprType      { return STRING }
func (node StringFunctionCall) Type() ExprType { return STRING }

func (node *ScalarLiteral) Eval(timestamp *time.Time) model.SampleValue {
	return node.value
}

func (node *ScalarArithExpr) Eval(timestamp *time.Time) model.SampleValue {
	lhs := node.lhs.Eval(timestamp)
	rhs := node.rhs.Eval(timestamp)
	return evalScalarBinop(node.opType, lhs, rhs)
}

func (node *ScalarFunctionCall) Eval(timestamp *time.Time) model.SampleValue {
	return node.function.callFn(timestamp, node.args).(model.SampleValue)
}

func (node *VectorAggregation) labelsToGroupingKey(labels model.Metric) string {
	keyParts := []string{}
	for _, keyLabel := range node.groupBy {
		keyParts = append(keyParts, string(labels[keyLabel]))
	}
	return strings.Join(keyParts, ",") // TODO not safe when label value contains comma.
}

func labelIntersection(metric1, metric2 model.Metric) model.Metric {
	intersection := model.Metric{}
	for label, value := range metric1 {
		if metric2[label] == value {
			intersection[label] = value
		}
	}
	return intersection
}

func (node *VectorAggregation) groupedAggregationsToVector(aggregations map[string]*groupedAggregation, timestamp *time.Time) Vector {
	vector := Vector{}
	for _, aggregation := range aggregations {
		if node.aggrType == AVG {
			aggregation.value = aggregation.value / model.SampleValue(aggregation.groupCount)
		}
		sample := &model.Sample{
			Metric:    aggregation.labels,
			Value:     aggregation.value,
			Timestamp: *timestamp,
		}
		vector = append(vector, sample)
	}
	return vector
}

func (node *VectorAggregation) Eval(timestamp *time.Time) Vector {
	vector := node.vector.Eval(timestamp)
	result := map[string]*groupedAggregation{}
	for _, sample := range vector {
		groupingKey := node.labelsToGroupingKey(sample.Metric)
		if groupedResult, ok := result[groupingKey]; ok {
			groupedResult.labels = labelIntersection(groupedResult.labels, sample.Metric)
			switch node.aggrType {
			case SUM:
				groupedResult.value += sample.Value
			case AVG:
				groupedResult.value += sample.Value
				groupedResult.groupCount++
			case MAX:
				if groupedResult.value < sample.Value {
					groupedResult.value = sample.Value
				}
			case MIN:
				if groupedResult.value > sample.Value {
					groupedResult.value = sample.Value
				}
			}
		} else {
			result[groupingKey] = &groupedAggregation{
				labels:     sample.Metric,
				value:      sample.Value,
				groupCount: 1,
			}
		}
	}
	return node.groupedAggregationsToVector(result, timestamp)
}

func (node *VectorLiteral) Eval(timestamp *time.Time) Vector {
	values, err := persistence.GetValueAtTime(node.labels, timestamp, &stalenessPolicy)
	if err != nil {
		log.Printf("Unable to get vector values")
		return Vector{}
	}
	return values
}

func (node *VectorFunctionCall) Eval(timestamp *time.Time) Vector {
	return node.function.callFn(timestamp, node.args).(Vector)
}

func evalScalarBinop(opType BinOpType,
	lhs model.SampleValue,
	rhs model.SampleValue) model.SampleValue {
	switch opType {
	case ADD:
		return lhs + rhs
	case SUB:
		return lhs - rhs
	case MUL:
		return lhs * rhs
	case DIV:
		if rhs != 0 {
			return lhs / rhs
		} else {
			return model.SampleValue(math.Inf(int(rhs)))
		}
	case MOD:
		if rhs != 0 {
			return model.SampleValue(int(lhs) % int(rhs))
		} else {
			return model.SampleValue(math.Inf(int(rhs)))
		}
	case EQ:
		if lhs == rhs {
			return 1
		} else {
			return 0
		}
	case NE:
		if lhs != rhs {
			return 1
		} else {
			return 0
		}
	case GT:
		if lhs > rhs {
			return 1
		} else {
			return 0
		}
	case LT:
		if lhs < rhs {
			return 1
		} else {
			return 0
		}
	case GE:
		if lhs >= rhs {
			return 1
		} else {
			return 0
		}
	case LE:
		if lhs <= rhs {
			return 1
		} else {
			return 0
		}
	}
	panic("Not all enum values enumerated in switch")
}

func evalVectorBinop(opType BinOpType,
	lhs model.SampleValue,
	rhs model.SampleValue) (model.SampleValue, bool) {
	switch opType {
	case ADD:
		return lhs + rhs, true
	case SUB:
		return lhs - rhs, true
	case MUL:
		return lhs * rhs, true
	case DIV:
		if rhs != 0 {
			return lhs / rhs, true
		} else {
			return model.SampleValue(math.Inf(int(rhs))), true
		}
	case MOD:
		if rhs != 0 {
			return model.SampleValue(int(lhs) % int(rhs)), true
		} else {
			return model.SampleValue(math.Inf(int(rhs))), true
		}
	case EQ:
		if lhs == rhs {
			return lhs, true
		} else {
			return 0, false
		}
	case NE:
		if lhs != rhs {
			return lhs, true
		} else {
			return 0, false
		}
	case GT:
		if lhs > rhs {
			return lhs, true
		} else {
			return 0, false
		}
	case LT:
		if lhs < rhs {
			return lhs, true
		} else {
			return 0, false
		}
	case GE:
		if lhs >= rhs {
			return lhs, true
		} else {
			return 0, false
		}
	case LE:
		if lhs <= rhs {
			return lhs, true
		} else {
			return 0, false
		}
	case AND:
		return lhs, true
	}
	panic("Not all enum values enumerated in switch")
}

func labelsEqual(labels1, labels2 model.Metric) bool {
	if len(labels1) != len(labels2) {
		return false
	}
	for label, value := range labels1 {
		if labels2[label] != value && label != "name" {
			return false
		}
	}
	return true
}

func (node *VectorArithExpr) Eval(timestamp *time.Time) Vector {
	lhs := node.lhs.Eval(timestamp)
	result := Vector{}
	if node.rhs.Type() == SCALAR {
		rhs := node.rhs.(ScalarNode).Eval(timestamp)
		for _, lhsSample := range lhs {
			value, keep := evalVectorBinop(node.opType, lhsSample.Value, rhs)
			if keep {
				lhsSample.Value = value
				result = append(result, lhsSample)
			}
		}
		return result
	} else if node.rhs.Type() == VECTOR {
		rhs := node.rhs.(VectorNode).Eval(timestamp)
		for _, lhsSample := range lhs {
			for _, rhsSample := range rhs {
				if labelsEqual(lhsSample.Metric, rhsSample.Metric) {
					value, keep := evalVectorBinop(node.opType, lhsSample.Value, rhsSample.Value)
					if keep {
						lhsSample.Value = value
						result = append(result, lhsSample)
					}
				}
			}
		}
		return result
	}
	panic("Invalid vector arithmetic expression operands")
}

func (node *MatrixLiteral) Eval(timestamp *time.Time) Matrix {
	interval := &model.Interval{
		OldestInclusive: timestamp.Add(-node.interval),
		NewestInclusive: *timestamp,
	}
	values, err := persistence.GetRangeValues(node.labels, interval, &stalenessPolicy)
	if err != nil {
		log.Printf("Unable to get values for vector interval")
		return Matrix{}
	}
	return values
}

func (node *MatrixLiteral) EvalBoundaries(timestamp *time.Time) Matrix {
	interval := &model.Interval{
		OldestInclusive: timestamp.Add(-node.interval),
		NewestInclusive: *timestamp,
	}
	values, err := persistence.GetBoundaryValues(node.labels, interval, &stalenessPolicy)
	if err != nil {
		log.Printf("Unable to get boundary values for vector interval")
		return Matrix{}
	}
	return values
}

func (node *StringLiteral) Eval(timestamp *time.Time) string {
	return node.str
}

func (node *StringFunctionCall) Eval(timestamp *time.Time) string {
	return node.function.callFn(timestamp, node.args).(string)
}

// ----------------------------------------------------------------------------
// Constructors.

func NewScalarLiteral(value model.SampleValue) *ScalarLiteral {
	return &ScalarLiteral{
		value: value,
	}
}

func NewVectorLiteral(labels model.LabelSet) *VectorLiteral {
	return &VectorLiteral{
		labels: labels,
	}
}

func NewVectorAggregation(aggrType AggrType, vector VectorNode, groupBy []model.LabelName) *VectorAggregation {
	return &VectorAggregation{
		aggrType: aggrType,
		groupBy:  groupBy,
		vector:   vector,
	}
}

func NewFunctionCall(function *Function, args []Node) (Node, error) {
	if err := function.CheckArgTypes(args); err != nil {
		return nil, err
	}
	switch function.returnType {
	case SCALAR:
		return &ScalarFunctionCall{
			function: function,
			args:     args,
		}, nil
	case VECTOR:
		return &VectorFunctionCall{
			function: function,
			args:     args,
		}, nil
	case STRING:
		return &StringFunctionCall{
			function: function,
			args:     args,
		}, nil
	}
	panic("Function with invalid return type")
}

func nodesHaveTypes(nodes []Node, exprTypes []ExprType) bool {
	for _, node := range nodes {
		for _, exprType := range exprTypes {
			if node.Type() == exprType {
				return true
			}
		}
	}
	return false
}

func NewArithExpr(opType BinOpType, lhs Node, rhs Node) (Node, error) {
	if !nodesHaveTypes([]Node{lhs, rhs}, []ExprType{SCALAR, VECTOR}) {
		return nil, errors.New("Binary operands must be of vector or scalar type")
	}
	if lhs.Type() == SCALAR && rhs.Type() == VECTOR {
		return nil, errors.New("Left side of vector binary operation must be of vector type")
	}

	if opType == AND || opType == OR {
		if lhs.Type() == SCALAR || rhs.Type() == SCALAR {
			return nil, errors.New("AND and OR operators may only be used between vectors")
		}
	}

	if lhs.Type() == VECTOR || rhs.Type() == VECTOR {
		return &VectorArithExpr{
			opType: opType,
			lhs:    lhs.(VectorNode),
			rhs:    rhs,
		}, nil
	}

	return &ScalarArithExpr{
		opType: opType,
		lhs:    lhs.(ScalarNode),
		rhs:    rhs.(ScalarNode),
	}, nil
}

func NewMatrixLiteral(vector *VectorLiteral, interval time.Duration) *MatrixLiteral {
	return &MatrixLiteral{
		labels:   vector.labels,
		interval: interval,
	}
}

func NewStringLiteral(str string) *StringLiteral {
	return &StringLiteral{
		str: str,
	}
}
<|endoftext|>


<commit_before>package ot

// np.go
// Naor-Pinkas oblivious transfer
//
// Efficient oblivious transfer protocols
// Moni Naor and Benny Pinkas
// SODA 2001
// https://dl.acm.org/citation.cfm?id=365502
//
// This is the "basic oblivious transfer protocol" of section 2.3

import (
	//	"log"
	"crypto/aes"
	"math/big"
)

const (
	KEY_SIZE     = aes.BlockSize
	HASH_SIZE    = 32
	MSG_LEN      = KEY_SIZE * 8
	NUM_PAIRS    = 1024 * 64 * 2
	NP_MSG_LEN   = NUM_PAIRS / 8
	SEC_PARAM    = 80
	primeHex     = "B10B8F96A080E01DDE92DE5EAE5D54EC52C99FBCFB06A3C69A6A9DCA52D23B616073E28675A23D189838EF1E2EE652C013ECB4AEA906112324975C3CD49B83BFACCBDD7D90C4BD7098488E9C219A73724EFFD6FAE5644738FAA31A4FF55BCCC0A151AF5F0DC8B4BD45BF37DF365C1A65E68CFDA76D4DA708DF1FB2BC2E4A4371"
	generatorHex = "A4D1CBD5C3FD34126765A442EFB99905F8104DD258AC507FD6406CFF14266D31266FEA1E5C41564B777E690F5504F213160217B4B01B886A5E91547F9E2749F4D7FBD7D3B9A92EE1909D0D2263F80A76A6A24C087A091F531DBF0A0169B6A28AD662A4D18E73AFA32D779D5918D08BC8858F4DCEF97C2A24855E6EEB22B3B2E5"
)

type HashedElGamalCiph struct {
	C0 big.Int
	C1 []byte
}

type NPSender struct {
	npSendPk   chan PublicKey
	npRecvPk   chan big.Int
	npSendEncs chan HashedElGamalCiph
	NPSenderParams
}

type NPSenderParams struct {
	pri    *PrivateKey
	MsgLen int
}

type NPReceiver struct {
	npSendPk   chan PublicKey
	npRecvPk   chan big.Int
	npSendEncs chan HashedElGamalCiph
	NPReceiverParams
}

type NPReceiverParams struct {
	PkSender PublicKey
	MsgLen   int
}

func GenNPParams(MsgLen int) (snd NPSenderParams, rcv NPReceiverParams) {
	pri := new(PrivateKey)
	pri.PublicKey = publicParams
	pri.X = generateNumNonce(publicParams.P)
	pri.Y = new(big.Int).Exp(pri.G, pri.X, pri.P)
	snd.pri = pri
	snd.MsgLen = MsgLen
	rcv.PkSender = pri.PublicKey
	rcv.MsgLen = MsgLen
	return
}

func NewNPSender(MsgLen int,
	snd NPSenderParams,
	npSendPk chan PublicKey,
	npRecvPk chan big.Int,
	npSendEncs chan HashedElGamalCiph) *NPSender {

	sender := new(NPSender)
	sender.npSendPk = npSendPk
	sender.npRecvPk = npRecvPk
	sender.npSendEncs = npSendEncs
	sender.NPSenderParams = snd

	return sender
}

func NewNPReceiver(MsgLen int, rcv NPReceiverParams,
	npSendPk chan PublicKey,
	npRecvPk chan big.Int,
	npSendEncs chan HashedElGamalCiph) *NPReceiver {

	receiver := new(NPReceiver)
	receiver.npSendPk = npSendPk
	receiver.npRecvPk = npRecvPk
	receiver.npSendEncs = npSendEncs
	receiver.NPReceiverParams = rcv

	return receiver
}

func NewNP(MsgLen int) (*NPSender, *NPReceiver) {
	snd, rcv := GenNPParams(MsgLen)
	npSendPk := make(chan PublicKey)
	npRecvPk := make(chan big.Int)
	npSendEncs := make(chan HashedElGamalCiph)

	return NewNPSender(MsgLen, snd, npSendPk, npRecvPk, npSendEncs),
		NewNPReceiver(MsgLen, rcv, npSendPk, npRecvPk, npSendEncs)
}

func (self *NPSender) Send(m0, m1 Message) {
	//	log.Println("Starting run of OTNPSender.Send")
	pks := make([]big.Int, 2)
	pks[0] = <-self.npRecvPk
	pks[1].ModInverse(&pks[0], publicParams.P)
	pks[1].Mul(&pks[1], self.pri.PublicKey.Y).Mod(&pks[1], publicParams.P)
	r0 := generateNumNonce(publicParams.P)
	r1 := generateNumNonce(publicParams.P)

	maskedVal0 := make([]byte, len(m0))
	xorBytes(maskedVal0, RO(expModP(&pks[0], r0).Bytes(), 8*len(m0)), m0)
	maskedVal1 := make([]byte, len(m0))
	xorBytes(maskedVal1, RO(expModP(&pks[1], r1).Bytes(), 8*len(m0)), m1)
	e0 := HashedElGamalCiph{C0: *gExpModP(r0), C1: maskedVal0}
	e1 := HashedElGamalCiph{C0: *gExpModP(r1), C1: maskedVal1}
	self.npSendEncs <- e0
	self.npSendEncs <- e1
	//	log.Println("Completed run of OTNPSender.Send")
}

func (self *NPReceiver) Receive(s Selector) Message {
	//	log.Printf("Starting run of OTNPSender.Receive. self.MsgLen=%d\n", self.MsgLen)
	pks := make([]*big.Int, 2)
	k := generateNumNonce(publicParams.P)
	pks[s] = gExpModP(k)
	pks[1-s] = new(big.Int).ModInverse(pks[s], publicParams.P)
	pks[1-s].Mul(pks[1-s], self.PkSender.Y).Mod(pks[1-s], publicParams.P)
	self.npRecvPk <- *pks[0]
	ciphs := []HashedElGamalCiph{<-self.npSendEncs, <-self.npSendEncs}
	res := make([]byte, self.MsgLen)
	xorBytes(res, RO(expModP(&ciphs[s].C0, k).Bytes(), 8*self.MsgLen), ciphs[s].C1)
	//	log.Printf("Completed run of OTNPSender.Receive. len(res)=%d\n", len(res))
	return res
}

func xorBytes(a, b, c []byte) {
	for i := range a {
		switch {
		case i < len(b) && i < len(c):
			a[i] = b[i] ^ c[i]
		case i < len(b):
			a[i] = b[i]
		case i < len(c):
			a[i] = c[i]
		}
	}
}

func expModP(g, x *big.Int) *big.Int {
	return new(big.Int).Exp(g, x, publicParams.P)
}

func gExpModP(x *big.Int) *big.Int {
	return expModP(publicParams.G, x)
}

var publicParams = PublicKey{G: fromHex(generatorHex), P: fromHex(primeHex)}

func fromHex(hex string) *big.Int {
	n, ok := new(big.Int).SetString(hex, 16)
	if !ok {
		panic("failed to parse hex number")
	}
	return n
}
<commit_msg>committing debugging changes for now<commit_after>package ot

// np.go
// Naor-Pinkas oblivious transfer
//
// Efficient oblivious transfer protocols
// Moni Naor and Benny Pinkas
// SODA 2001
// https://dl.acm.org/citation.cfm?id=365502
//
// This is the "basic oblivious transfer protocol" of section 2.3

import (
	//	"log"
	"crypto/aes"
	"math/big"
)

const (
	KEY_SIZE     = aes.BlockSize
	HASH_SIZE    = 32
	MSG_LEN      = KEY_SIZE * 8
	NUM_PAIRS    = 1024
	NP_MSG_LEN   = NUM_PAIRS / 8
	SEC_PARAM    = 80
	primeHex     = "B10B8F96A080E01DDE92DE5EAE5D54EC52C99FBCFB06A3C69A6A9DCA52D23B616073E28675A23D189838EF1E2EE652C013ECB4AEA906112324975C3CD49B83BFACCBDD7D90C4BD7098488E9C219A73724EFFD6FAE5644738FAA31A4FF55BCCC0A151AF5F0DC8B4BD45BF37DF365C1A65E68CFDA76D4DA708DF1FB2BC2E4A4371"
	generatorHex = "A4D1CBD5C3FD34126765A442EFB99905F8104DD258AC507FD6406CFF14266D31266FEA1E5C41564B777E690F5504F213160217B4B01B886A5E91547F9E2749F4D7FBD7D3B9A92EE1909D0D2263F80A76A6A24C087A091F531DBF0A0169B6A28AD662A4D18E73AFA32D779D5918D08BC8858F4DCEF97C2A24855E6EEB22B3B2E5"
)

type HashedElGamalCiph struct {
	C0 big.Int
	C1 []byte
}

type NPSender struct {
	npSendPk   chan PublicKey
	npRecvPk   chan big.Int
	npSendEncs chan HashedElGamalCiph
	NPSenderParams
}

type NPSenderParams struct {
	pri    *PrivateKey
	MsgLen int
}

type NPReceiver struct {
	npSendPk   chan PublicKey
	npRecvPk   chan big.Int
	npSendEncs chan HashedElGamalCiph
	NPReceiverParams
}

type NPReceiverParams struct {
	PkSender PublicKey
	MsgLen   int
}

func GenNPParams(MsgLen int) (snd NPSenderParams, rcv NPReceiverParams) {
	pri := new(PrivateKey)
	pri.PublicKey = publicParams
	pri.X = generateNumNonce(publicParams.P)
	pri.Y = new(big.Int).Exp(pri.G, pri.X, pri.P)
	snd.pri = pri
	snd.MsgLen = MsgLen
	rcv.PkSender = pri.PublicKey
	rcv.MsgLen = MsgLen
	return
}

func NewNPSender(MsgLen int,
	snd NPSenderParams,
	npSendPk chan PublicKey,
	npRecvPk chan big.Int,
	npSendEncs chan HashedElGamalCiph) *NPSender {

	sender := new(NPSender)
	sender.npSendPk = npSendPk
	sender.npRecvPk = npRecvPk
	sender.npSendEncs = npSendEncs
	sender.NPSenderParams = snd

	return sender
}

func NewNPReceiver(MsgLen int, rcv NPReceiverParams,
	npSendPk chan PublicKey,
	npRecvPk chan big.Int,
	npSendEncs chan HashedElGamalCiph) *NPReceiver {

	receiver := new(NPReceiver)
	receiver.npSendPk = npSendPk
	receiver.npRecvPk = npRecvPk
	receiver.npSendEncs = npSendEncs
	receiver.NPReceiverParams = rcv

	return receiver
}

func NewNP(MsgLen int) (*NPSender, *NPReceiver) {
	snd, rcv := GenNPParams(MsgLen)
	npSendPk := make(chan PublicKey)
	npRecvPk := make(chan big.Int)
	npSendEncs := make(chan HashedElGamalCiph)

	return NewNPSender(MsgLen, snd, npSendPk, npRecvPk, npSendEncs),
		NewNPReceiver(MsgLen, rcv, npSendPk, npRecvPk, npSendEncs)
}

func (self *NPSender) Send(m0, m1 Message) {
	//	log.Println("Starting run of OTNPSender.Send")
	pks := make([]big.Int, 2)
	pks[0] = <-self.npRecvPk
	pks[1].ModInverse(&pks[0], publicParams.P)
	pks[1].Mul(&pks[1], self.pri.PublicKey.Y).Mod(&pks[1], publicParams.P)
	r0 := generateNumNonce(publicParams.P)
	r1 := generateNumNonce(publicParams.P)

	maskedVal0 := make([]byte, len(m0))
	xorBytes(maskedVal0, RO(expModP(&pks[0], r0).Bytes(), 8*len(m0)), m0)
	maskedVal1 := make([]byte, len(m0))
	xorBytes(maskedVal1, RO(expModP(&pks[1], r1).Bytes(), 8*len(m0)), m1)
	e0 := HashedElGamalCiph{C0: *gExpModP(r0), C1: maskedVal0}
	e1 := HashedElGamalCiph{C0: *gExpModP(r1), C1: maskedVal1}
	self.npSendEncs <- e0
	self.npSendEncs <- e1
	//	log.Println("Completed run of OTNPSender.Send")
}

func (self *NPReceiver) Receive(s Selector) Message {
	//	log.Printf("Starting run of OTNPSender.Receive. self.MsgLen=%d\n", self.MsgLen)
	pks := make([]*big.Int, 2)
	k := generateNumNonce(publicParams.P)
	pks[s] = gExpModP(k)
	pks[1-s] = new(big.Int).ModInverse(pks[s], publicParams.P)
	pks[1-s].Mul(pks[1-s], self.PkSender.Y).Mod(pks[1-s], publicParams.P)
	self.npRecvPk <- *pks[0]
	ciphs := []HashedElGamalCiph{<-self.npSendEncs, <-self.npSendEncs}
	res := make([]byte, self.MsgLen)
	xorBytes(res, RO(expModP(&ciphs[s].C0, k).Bytes(), 8*self.MsgLen), ciphs[s].C1)
	//	log.Printf("Completed run of OTNPSender.Receive. len(res)=%d\n", len(res))
	return res
}

func xorBytes(a, b, c []byte) {
	for i := range a {
		switch {
		case i < len(b) && i < len(c):
			a[i] = b[i] ^ c[i]
		case i < len(b):
			a[i] = b[i]
		case i < len(c):
			a[i] = c[i]
		}
	}
}

func expModP(g, x *big.Int) *big.Int {
	return new(big.Int).Exp(g, x, publicParams.P)
}

func gExpModP(x *big.Int) *big.Int {
	return expModP(publicParams.G, x)
}

var publicParams = PublicKey{G: fromHex(generatorHex), P: fromHex(primeHex)}

func fromHex(hex string) *big.Int {
	n, ok := new(big.Int).SetString(hex, 16)
	if !ok {
		panic("failed to parse hex number")
	}
	return n
}
<|endoftext|>