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

// Usage:
//
// import "github.com/astaxie/beego/logs"
//
//	log := NewLogger(10000)
//	log.SetLogger("console", "")
//
//	> the first params stand for how many channel
//
// Use it like this:
//
//	log.Trace("trace")
//	log.Info("info")
//	log.Warn("warning")
//	log.Debug("debug")
//	log.Critical("critical")
//
//  more docs http://beego.me/docs/module/logs.md
package logs

import (
	"encoding/json"
	"fmt"
	"os/exec"
	"path"
	"reflect"
	"runtime"
	"strings"
	"sync"

	"github.com/gogap/errors"
)

// RFC5424 log message levels.
const (
	LevelError = iota
	LevelWarn
	LevelInfo
	LevelDebug
)

// Legacy loglevel constants to ensure backwards compatibility.
//
// Deprecated: will be removed in 1.5.0.

type loggerType func() LoggerInterface

// LoggerInterface defines the behavior of a log provider.
type LoggerInterface interface {
	Init(config string) error
	WriteMsg(msg string, level int) error
	Destroy()
	Flush()
}

var adapters = make(map[string]loggerType)

// Register makes a log provide available by the provided name.
// If Register is called twice with the same name or if driver is nil,
// it panics.
func Register(name string, log loggerType) {
	if log == nil {
		panic("logs: Register provide is nil")
	}
	if _, dup := adapters[name]; dup {
		panic("logs: Register called twice for provider " + name)
	}
	adapters[name] = log
}

// BeeLogger is default logger in beego application.
// it can contain several providers and log message into all providers.
type BeeLogger struct {
	lock                sync.Mutex
	level               int
	enableFuncCallDepth bool
	loggerFuncCallDepth int
	msg                 chan *logMsg
	outputs             map[string]LoggerInterface
}

type logMsg struct {
	level int
	msg   string
}

// NewLogger returns a new BeeLogger.
// channellen means the number of messages in chan.
// if the buffering chan is full, logger adapters write to file or other way.
func NewLogger(channellen int64) *BeeLogger {
	bl := new(BeeLogger)
	bl.level = LevelDebug
	bl.loggerFuncCallDepth = 2
	bl.msg = make(chan *logMsg, channellen)
	bl.outputs = make(map[string]LoggerInterface)
	//bl.SetLogger("console", "") // default output to console
	go bl.startLogger()
	return bl
}

func NewFileLogger(file string) *BeeLogger {
	l := NewLogger(1024)
	path := strings.Split(file, "/")
	if len(path) > 1 {
		exec.Command("mkdir", path[0]).Run()
	}
	l.SetLogger("console", "")
	l.SetLogger("file", fmt.Sprintf(`{"filename":"%s","maxdays":7}`, file))
	l.EnableFuncCallDepth(true)
	l.SetLogFuncCallDepth(2)
	return l
}

// SetLogger provides a given logger adapter into BeeLogger with config string.
// config need to be correct JSON as string: {"interval":360}.
func (bl *BeeLogger) SetLogger(adaptername string, config string) error {
	bl.lock.Lock()
	defer bl.lock.Unlock()
	if log, ok := adapters[adaptername]; ok {
		lg := log()
		err := lg.Init(config)
		bl.outputs[adaptername] = lg
		if err != nil {
			fmt.Println("logs.BeeLogger.SetLogger: " + err.Error())
			return err
		}
	} else {
		return fmt.Errorf("logs: unknown adaptername %q (forgotten Register?)", adaptername)
	}
	return nil
}

// remove a logger adapter in BeeLogger.
func (bl *BeeLogger) DelLogger(adaptername string) error {
	bl.lock.Lock()
	defer bl.lock.Unlock()
	if lg, ok := bl.outputs[adaptername]; ok {
		lg.Destroy()
		delete(bl.outputs, adaptername)
		return nil
	} else {
		return fmt.Errorf("logs: unknown adaptername %q (forgotten Register?)", adaptername)
	}
}

func (bl *BeeLogger) writerMsg(loglevel int, msg string) error {
	if loglevel > bl.level {
		return nil
	}
	lm := new(logMsg)
	lm.level = loglevel
	if bl.enableFuncCallDepth {
		_, file, line, ok := runtime.Caller(bl.loggerFuncCallDepth)
		if ok {
			_, filename := path.Split(file)
			lm.msg = fmt.Sprintf("[%s:%d] %s", filename, line, msg)
		} else {
			lm.msg = msg
		}
	} else {
		lm.msg = msg
	}
	bl.msg <- lm
	return nil
}

// Set log message level.
//
// If message level (such as LevelDebug) is higher than logger level (such as LevelWarning),
// log providers will not even be sent the message.
func (bl *BeeLogger) SetLevel(l int) {
	bl.level = l
}

// set log funcCallDepth
func (bl *BeeLogger) SetLogFuncCallDepth(d int) {
	bl.loggerFuncCallDepth = d
}

// enable log funcCallDepth
func (bl *BeeLogger) EnableFuncCallDepth(b bool) {
	bl.enableFuncCallDepth = b
}

// start logger chan reading.
// when chan is not empty, write logs.
func (bl *BeeLogger) startLogger() {
	for {
		select {
		case bm := <-bl.msg:
			for _, l := range bl.outputs {
				err := l.WriteMsg(bm.msg, bm.level)
				if err != nil {
					fmt.Println("ERROR, unable to WriteMsg:", err)
				}
			}
		}
	}
}

// Log ERROR level message.
func (bl *BeeLogger) Error(v ...interface{}) {
	bl.log("Error", LevelError, v)
}

// Log WARNING level message.
func (bl *BeeLogger) Warn(v ...interface{}) {
	bl.log("Warn", LevelWarn, v)
}

// Log INFORMATIONAL level message.
func (bl *BeeLogger) Info(v ...interface{}) {
	bl.log("Info", LevelInfo, v)
}

// Log DEBUG level message.
func (bl *BeeLogger) Debug(v ...interface{}) {
	bl.log("Debug", LevelDebug, v)
}

func (bl *BeeLogger) log(tp string, level int, v ...interface{}) {
	msg := fmt.Sprintf("["+tp+"] "+generateFmtStr(len(v)), v...)
	for _, item := range v {
		if items, ok := item.([]interface{}); ok {
			if len(items) > 0 {
				item = items[0]
			}
		}
		stack := handleError(item)
		if stack != "" {
			msg = msg + "\n" + stack
		}
	}
	bl.writerMsg(level, msg)
}

func (bl *BeeLogger) Pretty(v interface{}, message string) {
	b, _ := json.MarshalIndent(v, " ", "  ")
	if message == "" {
		message = reflect.TypeOf(v).String()
	}
	bl.writerMsg(LevelDebug, fmt.Sprintf("[Pretty]\n%s\n%s", message, string(b)))
}

// flush all chan data.
func (bl *BeeLogger) Flush() {
	for _, l := range bl.outputs {
		l.Flush()
	}
}

// close logger, flush all chan data and destroy all adapters in BeeLogger.
func (bl *BeeLogger) Close() {
	for {
		if len(bl.msg) > 0 {
			bm := <-bl.msg
			for _, l := range bl.outputs {
				err := l.WriteMsg(bm.msg, bm.level)
				if err != nil {
					fmt.Println("ERROR, unable to WriteMsg (while closing logger):", err)
				}
			}
		} else {
			break
		}
	}
	for _, l := range bl.outputs {
		l.Flush()
		l.Destroy()
	}
}

func generateFmtStr(n int) string {
	return strings.Repeat("%v ", n)
}

func handleError(v interface{}) (msg string) {
	if err, ok := v.(errors.ErrCode); ok {
		msg = msg + err.StackTrace()
	}
	return
}
<commit_msg>fix console<commit_after>// Copyright 2014 beego Author. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Usage:
//
// import "github.com/astaxie/beego/logs"
//
//	log := NewLogger(10000)
//	log.SetLogger("console", "")
//
//	> the first params stand for how many channel
//
// Use it like this:
//
//	log.Trace("trace")
//	log.Info("info")
//	log.Warn("warning")
//	log.Debug("debug")
//	log.Critical("critical")
//
//  more docs http://beego.me/docs/module/logs.md
package logs

import (
	"encoding/json"
	"fmt"
	"os/exec"
	"path"
	"reflect"
	"runtime"
	"strings"
	"sync"

	"github.com/gogap/errors"
)

// RFC5424 log message levels.
const (
	LevelError = iota
	LevelWarn
	LevelInfo
	LevelDebug
)

// Legacy loglevel constants to ensure backwards compatibility.
//
// Deprecated: will be removed in 1.5.0.

type loggerType func() LoggerInterface

// LoggerInterface defines the behavior of a log provider.
type LoggerInterface interface {
	Init(config string) error
	WriteMsg(msg string, level int) error
	Destroy()
	Flush()
}

var adapters = make(map[string]loggerType)

// Register makes a log provide available by the provided name.
// If Register is called twice with the same name or if driver is nil,
// it panics.
func Register(name string, log loggerType) {
	if log == nil {
		panic("logs: Register provide is nil")
	}
	if _, dup := adapters[name]; dup {
		panic("logs: Register called twice for provider " + name)
	}
	adapters[name] = log
}

// BeeLogger is default logger in beego application.
// it can contain several providers and log message into all providers.
type BeeLogger struct {
	lock                sync.Mutex
	level               int
	enableFuncCallDepth bool
	loggerFuncCallDepth int
	msg                 chan *logMsg
	outputs             map[string]LoggerInterface
}

type logMsg struct {
	level int
	msg   string
}

// NewLogger returns a new BeeLogger.
// channellen means the number of messages in chan.
// if the buffering chan is full, logger adapters write to file or other way.
func NewLogger(channellen int64) *BeeLogger {
	bl := new(BeeLogger)
	bl.level = LevelDebug
	bl.loggerFuncCallDepth = 2
	bl.msg = make(chan *logMsg, channellen)
	bl.outputs = make(map[string]LoggerInterface)
	bl.SetLogger("console", "") // default output to console
	go bl.startLogger()
	return bl
}

func NewFileLogger(file string) *BeeLogger {
	l := NewLogger(1024)
	path := strings.Split(file, "/")
	if len(path) > 1 {
		exec.Command("mkdir", path[0]).Run()
	}
	l.SetLogger("file", fmt.Sprintf(`{"filename":"%s","maxdays":7}`, file))
	l.EnableFuncCallDepth(true)
	l.SetLogFuncCallDepth(2)
	return l
}

// SetLogger provides a given logger adapter into BeeLogger with config string.
// config need to be correct JSON as string: {"interval":360}.
func (bl *BeeLogger) SetLogger(adaptername string, config string) error {
	bl.lock.Lock()
	defer bl.lock.Unlock()
	if log, ok := adapters[adaptername]; ok {
		lg := log()
		err := lg.Init(config)
		bl.outputs[adaptername] = lg
		if err != nil {
			fmt.Println("logs.BeeLogger.SetLogger: " + err.Error())
			return err
		}
	} else {
		return fmt.Errorf("logs: unknown adaptername %q (forgotten Register?)", adaptername)
	}
	return nil
}

// remove a logger adapter in BeeLogger.
func (bl *BeeLogger) DelLogger(adaptername string) error {
	bl.lock.Lock()
	defer bl.lock.Unlock()
	if lg, ok := bl.outputs[adaptername]; ok {
		lg.Destroy()
		delete(bl.outputs, adaptername)
		return nil
	} else {
		return fmt.Errorf("logs: unknown adaptername %q (forgotten Register?)", adaptername)
	}
}

func (bl *BeeLogger) writerMsg(loglevel int, msg string) error {
	if loglevel > bl.level {
		return nil
	}
	lm := new(logMsg)
	lm.level = loglevel
	if bl.enableFuncCallDepth {
		_, file, line, ok := runtime.Caller(bl.loggerFuncCallDepth)
		if ok {
			_, filename := path.Split(file)
			lm.msg = fmt.Sprintf("[%s:%d] %s", filename, line, msg)
		} else {
			lm.msg = msg
		}
	} else {
		lm.msg = msg
	}
	bl.msg <- lm
	return nil
}

// Set log message level.
//
// If message level (such as LevelDebug) is higher than logger level (such as LevelWarning),
// log providers will not even be sent the message.
func (bl *BeeLogger) SetLevel(l int) {
	bl.level = l
}

// set log funcCallDepth
func (bl *BeeLogger) SetLogFuncCallDepth(d int) {
	bl.loggerFuncCallDepth = d
}

// enable log funcCallDepth
func (bl *BeeLogger) EnableFuncCallDepth(b bool) {
	bl.enableFuncCallDepth = b
}

// start logger chan reading.
// when chan is not empty, write logs.
func (bl *BeeLogger) startLogger() {
	for {
		select {
		case bm := <-bl.msg:
			for _, l := range bl.outputs {
				err := l.WriteMsg(bm.msg, bm.level)
				if err != nil {
					fmt.Println("ERROR, unable to WriteMsg:", err)
				}
			}
		}
	}
}

// Log ERROR level message.
func (bl *BeeLogger) Error(v ...interface{}) {
	bl.log("Error", LevelError, v)
}

// Log WARNING level message.
func (bl *BeeLogger) Warn(v ...interface{}) {
	bl.log("Warn", LevelWarn, v)
}

// Log INFORMATIONAL level message.
func (bl *BeeLogger) Info(v ...interface{}) {
	bl.log("Info", LevelInfo, v)
}

// Log DEBUG level message.
func (bl *BeeLogger) Debug(v ...interface{}) {
	bl.log("Debug", LevelDebug, v)
}

func (bl *BeeLogger) log(tp string, level int, v ...interface{}) {
	msg := fmt.Sprintf("["+tp+"] "+generateFmtStr(len(v)), v...)
	for _, item := range v {
		if items, ok := item.([]interface{}); ok {
			if len(items) > 0 {
				item = items[0]
			}
		}
		stack := handleError(item)
		if stack != "" {
			msg = msg + "\n" + stack
		}
	}
	bl.writerMsg(level, msg)
}

func (bl *BeeLogger) Pretty(v interface{}, message string) {
	b, _ := json.MarshalIndent(v, " ", "  ")
	if message == "" {
		message = reflect.TypeOf(v).String()
	}
	bl.writerMsg(LevelDebug, fmt.Sprintf("[Pretty]\n%s\n%s", message, string(b)))
}

// flush all chan data.
func (bl *BeeLogger) Flush() {
	for _, l := range bl.outputs {
		l.Flush()
	}
}

// close logger, flush all chan data and destroy all adapters in BeeLogger.
func (bl *BeeLogger) Close() {
	for {
		if len(bl.msg) > 0 {
			bm := <-bl.msg
			for _, l := range bl.outputs {
				err := l.WriteMsg(bm.msg, bm.level)
				if err != nil {
					fmt.Println("ERROR, unable to WriteMsg (while closing logger):", err)
				}
			}
		} else {
			break
		}
	}
	for _, l := range bl.outputs {
		l.Flush()
		l.Destroy()
	}
}

func generateFmtStr(n int) string {
	return strings.Repeat("%v ", n)
}

func handleError(v interface{}) (msg string) {
	if err, ok := v.(errors.ErrCode); ok {
		msg = msg + err.StackTrace()
	}
	return
}
<|endoftext|>


<commit_before>//  Copyright (c) 2014 Couchbase, Inc.

package gofast

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

// Logger interface for gofast logging, applications can
// supply a logger object implementing this interface or
// gofast will fall back to the defaul logger.
type Logger interface {
	Fatalf(format string, v ...interface{})
	Errorf(format string, v ...interface{})
	Warnf(format string, v ...interface{})
	Infof(format string, v ...interface{})
	Verbosef(format string, v ...interface{})
	Debugf(format string, v ...interface{})
	Tracef(format string, v ...interface{})
}

// * default log file is os.Stdout
// * default level is LogLevelInfo
type DefaultLogger struct {
	level  logLevel
	output io.Writer
}

type logLevel int

const (
	logLevelIgnore logLevel = iota + 1
	logLevelFatal
	logLevelError
	logLevelWarn
	logLevelInfo
	logLevelVerbose
	logLevelDebug
	logLevelTrace
)

var log Logger // object used by gofast component for logging.

func setLogger(logger Logger, config map[string]interface{}) Logger {
	if logger != nil {
		log = logger
		return log
	}

	var err error
	level := logLevelInfo
	if val, ok := config["log.level"]; ok {
		level = string2logLevel(val.(string))
	}
	logfd := os.Stdout
	if val, ok := config["log.file"]; ok {
		logfile := val.(string)
		logfd, err = os.OpenFile(logfile, os.O_RDWR|os.O_APPEND, 0660)
		if err != nil {
			if logfd, err = os.Create(logfile); err != nil {
				panic(err)
			}
		}
	}
	log = &DefaultLogger{level: level, output: logfd}
	return log
}

func (l *DefaultLogger) Fatalf(format string, v ...interface{}) {
	l.printf(logLevelFatal, format, v...)
}

func (l *DefaultLogger) Errorf(format string, v ...interface{}) {
	l.printf(logLevelError, format, v...)
}

func (l *DefaultLogger) Warnf(format string, v ...interface{}) {
	l.printf(logLevelWarn, format, v...)
}

func (l *DefaultLogger) Infof(format string, v ...interface{}) {
	l.printf(logLevelInfo, format, v...)
}

func (l *DefaultLogger) Verbosef(format string, v ...interface{}) {
	l.printf(logLevelVerbose, format, v...)
}

func (l *DefaultLogger) Debugf(format string, v ...interface{}) {
	l.printf(logLevelDebug, format, v...)
}

func (l *DefaultLogger) Tracef(format string, v ...interface{}) {
	l.printf(logLevelTrace, format, v...)
}

func (l *DefaultLogger) printf(level logLevel, format string, v ...interface{}) {
	if l.canlog(level) {
		ts := time.Now().Format("2006-01-02T15:04:05.999Z-07:00")
		fmt.Fprintf(l.output, ts+" ["+l.level.String()+"] "+format, v...)
	}
}

func (l *DefaultLogger) canlog(level logLevel) bool {
	if level <= l.level {
		return true
	}
	return false
}

func (l logLevel) String() string {
	switch l {
	case logLevelIgnore:
		return "Ignore"
	case logLevelFatal:
		return "Fatal"
	case logLevelError:
		return "Error"
	case logLevelWarn:
		return "Warn"
	case logLevelInfo:
		return "Info"
	case logLevelVerbose:
		return "Verbose"
	case logLevelDebug:
		return "Debug"
	case logLevelTrace:
		return "Trace"
	}
	panic("unexpected log level") // should never reach here
}

func string2logLevel(s string) logLevel {
	s = strings.ToLower(s)
	switch s {
	case "ignore":
		return logLevelIgnore
	case "fatal":
		return logLevelFatal
	case "error":
		return logLevelError
	case "warn":
		return logLevelWarn
	case "info":
		return logLevelInfo
	case "verbose":
		return logLevelVerbose
	case "debug":
		return logLevelDebug
	case "trace":
		return logLevelTrace
	}
	panic("unexpected log level") // should never reach here
}
<commit_msg>bug fix in log leveling.<commit_after>//  Copyright (c) 2014 Couchbase, Inc.

package gofast

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

// Logger interface for gofast logging, applications can
// supply a logger object implementing this interface or
// gofast will fall back to the defaul logger.
type Logger interface {
	Fatalf(format string, v ...interface{})
	Errorf(format string, v ...interface{})
	Warnf(format string, v ...interface{})
	Infof(format string, v ...interface{})
	Verbosef(format string, v ...interface{})
	Debugf(format string, v ...interface{})
	Tracef(format string, v ...interface{})
}

// * default log file is os.Stdout
// * default level is LogLevelInfo
type DefaultLogger struct {
	level  logLevel
	output io.Writer
}

type logLevel int

const (
	logLevelIgnore logLevel = iota + 1
	logLevelFatal
	logLevelError
	logLevelWarn
	logLevelInfo
	logLevelVerbose
	logLevelDebug
	logLevelTrace
)

var log Logger // object used by gofast component for logging.

func setLogger(logger Logger, config map[string]interface{}) Logger {
	if logger != nil {
		log = logger
		return log
	}

	var err error
	level := logLevelInfo
	if val, ok := config["log.level"]; ok {
		level = string2logLevel(val.(string))
	}
	logfd := os.Stdout
	if val, ok := config["log.file"]; ok {
		logfile := val.(string)
		logfd, err = os.OpenFile(logfile, os.O_RDWR|os.O_APPEND, 0660)
		if err != nil {
			if logfd, err = os.Create(logfile); err != nil {
				panic(err)
			}
		}
	}
	log = &DefaultLogger{level: level, output: logfd}
	return log
}

func (l *DefaultLogger) Fatalf(format string, v ...interface{}) {
	l.printf(logLevelFatal, format, v...)
}

func (l *DefaultLogger) Errorf(format string, v ...interface{}) {
	l.printf(logLevelError, format, v...)
}

func (l *DefaultLogger) Warnf(format string, v ...interface{}) {
	l.printf(logLevelWarn, format, v...)
}

func (l *DefaultLogger) Infof(format string, v ...interface{}) {
	l.printf(logLevelInfo, format, v...)
}

func (l *DefaultLogger) Verbosef(format string, v ...interface{}) {
	l.printf(logLevelVerbose, format, v...)
}

func (l *DefaultLogger) Debugf(format string, v ...interface{}) {
	l.printf(logLevelDebug, format, v...)
}

func (l *DefaultLogger) Tracef(format string, v ...interface{}) {
	l.printf(logLevelTrace, format, v...)
}

func (l *DefaultLogger) printf(level logLevel, format string, v ...interface{}) {
	if l.canlog(level) {
		ts := time.Now().Format("2006-01-02T15:04:05.999Z-07:00")
		fmt.Fprintf(l.output, ts+" ["+level.String()+"] "+format, v...)
	}
}

func (l *DefaultLogger) canlog(level logLevel) bool {
	if level <= l.level {
		return true
	}
	return false
}

func (l logLevel) String() string {
	switch l {
	case logLevelIgnore:
		return "Ignore"
	case logLevelFatal:
		return "Fatal"
	case logLevelError:
		return "Error"
	case logLevelWarn:
		return "Warn"
	case logLevelInfo:
		return "Info"
	case logLevelVerbose:
		return "Verbose"
	case logLevelDebug:
		return "Debug"
	case logLevelTrace:
		return "Trace"
	}
	panic("unexpected log level") // should never reach here
}

func string2logLevel(s string) logLevel {
	s = strings.ToLower(s)
	switch s {
	case "ignore":
		return logLevelIgnore
	case "fatal":
		return logLevelFatal
	case "error":
		return logLevelError
	case "warn":
		return logLevelWarn
	case "info":
		return logLevelInfo
	case "verbose":
		return logLevelVerbose
	case "debug":
		return logLevelDebug
	case "trace":
		return logLevelTrace
	}
	panic("unexpected log level") // should never reach here
}
<|endoftext|>


<commit_before>package main

import "fmt"
import "io"
import "os"
import "os/exec"
import "strings"

import . "./lua"
import "./alias"
import "./conio/readline"
import "./dos"
import "./interpreter"
import "./mbcs"

import "github.com/shiena/ansicolor"

type LuaFunction struct {
	L            *Lua
	registoryKey string
}

var LuaInstanceToCmd = map[uintptr]*interpreter.Interpreter{}

func (this LuaFunction) String() string {
	return "<<Lua-function>>"
}

func (this LuaFunction) Call(cmd *interpreter.Interpreter) (interpreter.NextT, error) {
	this.L.GetField(Registory, this.registoryKey)
	this.L.NewTable()
	for i, arg1 := range cmd.Args {
		this.L.PushInteger(i)
		this.L.PushString(arg1)
		this.L.SetTable(-3)
	}
	LuaInstanceToCmd[this.L.Id()] = cmd
	err := this.L.Call(1, 0)
	return interpreter.CONTINUE, err
}

func cmdAlias(L *Lua) int {
	name, nameErr := L.ToString(1)
	if nameErr != nil {
		L.PushNil()
		L.PushString(nameErr.Error())
		return 2
	}
	key := strings.ToLower(name)
	switch L.GetType(2) {
	case TSTRING:
		value, err := L.ToString(2)
		if err == nil {
			alias.Table[key] = alias.New(value)
		} else {
			L.PushNil()
			L.PushString(err.Error())
			return 2
		}
	case TFUNCTION:
		regkey := "nyagos.alias." + key
		L.SetField(Registory, regkey)
		alias.Table[key] = LuaFunction{L, regkey}
	}
	L.PushBool(true)
	return 1
}

func cmdSetEnv(L *Lua) int {
	name, nameErr := L.ToString(1)
	if nameErr != nil {
		L.PushNil()
		L.PushString(nameErr.Error())
		return 2
	}
	value, valueErr := L.ToString(2)
	if valueErr != nil {
		L.PushNil()
		L.PushString(valueErr.Error())
		return 2
	}
	os.Setenv(name, value)
	L.PushBool(true)
	return 1
}

func cmdGetEnv(L *Lua) int {
	name, nameErr := L.ToString(1)
	if nameErr != nil {
		L.PushNil()
		return 1
	}
	value := os.Getenv(name)
	if len(value) > 0 {
		L.PushString(value)
	} else {
		L.PushNil()
	}
	return 1
}

func cmdExec(L *Lua) int {
	statement, statementErr := L.ToString(1)
	if statementErr != nil {
		L.PushNil()
		L.PushString(statementErr.Error())
		return 2
	}
	_, err := interpreter.New().Interpret(statement)

	if err != nil {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
	L.PushBool(true)
	return 1
}

func cmdEval(L *Lua) int {
	statement, statementErr := L.ToString(1)
	if statementErr != nil {
		L.PushNil()
		L.PushString(statementErr.Error())
		return 2
	}
	r, w, err := os.Pipe()
	if err != nil {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
	go func(statement string, w *os.File) {
		it := interpreter.New()
		it.Stdout = w
		it.Interpret(statement)
		w.Close()
	}(statement, w)

	var result = []byte{}
	for {
		buffer := make([]byte, 256)
		size, err := r.Read(buffer)
		if err != nil || size <= 0 {
			break
		}
		result = append(result, buffer[0:size]...)
	}
	r.Close()
	L.PushAnsiString(result)
	return 1
}

func cmdWrite(L *Lua) int {
	var out io.Writer = os.Stdout
	cmd, cmdOk := LuaInstanceToCmd[L.Id()]
	if cmdOk && cmd != nil && cmd.Stdout != nil {
		out = cmd.Stdout
	}
	switch out.(type) {
	case *os.File:
		out = ansicolor.NewAnsiColorWriter(out)
	}

	n := L.GetTop()
	for i := 1; i <= n; i++ {
		str, err := L.ToString(i)
		if err != nil {
			L.PushNil()
			L.PushString(err.Error())
			return 2
		}
		if i > 1 {
			fmt.Fprint(out, "\t")
		}
		fmt.Fprint(out, str)
	}
	L.PushBool(true)
	return 1
}

func cmdGetwd(L *Lua) int {
	wd, err := os.Getwd()
	if err == nil {
		L.PushString(wd)
		return 1
	} else {
		return 0
	}
}

func cmdWhich(L *Lua) int {
	if L.GetType(-1) != TSTRING {
		return 0
	}
	name, nameErr := L.ToString(-1)
	if nameErr != nil {
		L.PushNil()
		L.PushString(nameErr.Error())
		return 2
	}
	path, err := exec.LookPath(name)
	if err == nil {
		L.PushString(path)
		return 1
	} else {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
}

func cmdAtoU(L *Lua) int {
	str, err := mbcs.AtoU(L.ToAnsiString(1))
	if err == nil {
		L.PushString(str)
		return 1
	} else {
		return 0
	}
}

func cmdUtoA(L *Lua) int {
	utf8, utf8err := L.ToString(1)
	if utf8err != nil {
		L.PushNil()
		L.PushString(utf8err.Error())
		return 2
	}
	str, err := mbcs.UtoA(utf8)
	if err == nil {
		if len(str) >= 1 {
			L.PushAnsiString(str[:len(str)-1])
		} else {
			L.PushString("")
		}
		L.PushNil()
		return 2
	} else {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
}

func cmdGlob(L *Lua) int {
	if !L.IsString(-1) {
		return 0
	}
	wildcard, wildcardErr := L.ToString(-1)
	if wildcardErr != nil {
		L.PushNil()
		L.PushString(wildcardErr.Error())
		return 2
	}
	list, err := dos.Glob(wildcard)
	if err != nil {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	} else {
		L.NewTable()
		for i := 0; i < len(list); i++ {
			L.PushInteger(i + 1)
			L.PushString(list[i])
			L.SetTable(-3)
		}
		return 1
	}
}

type KeyLuaFuncT struct {
	L            *Lua
	registoryKey string
}

func (this *KeyLuaFuncT) Call(buffer *readline.Buffer) readline.Result {
	this.L.GetField(Registory, this.registoryKey)
	if err := this.L.Call(0, 0); err != nil {
		fmt.Println(os.Stderr, err)
	}
	return readline.CONTINUE
}

func cmdBindKey(L *Lua) int {
	key, keyErr := L.ToString(-2)
	if keyErr != nil {
		L.PushString(keyErr.Error())
		return 1
	}
	key = strings.Replace(strings.ToUpper(key), "-", "_", -1)
	switch L.GetType(-1) {
	case TFUNCTION:
		regkey := "nyagos.bind." + key
		L.SetField(Registory, regkey)
		if err := readline.BindKeyFunc(key, &KeyLuaFuncT{L, regkey}); err != nil {
			L.PushNil()
			L.PushString(err.Error())
			return 2
		} else {
			L.PushBool(true)
			return 1
		}
	default:
		val, valErr := L.ToString(-1)
		if valErr != nil {
			L.PushNil()
			L.PushString(valErr.Error())
			return 2
		}
		err := readline.BindKeySymbol(key, val)
		if err != nil {
			L.PushNil()
			L.PushString(err.Error())
			return 2
		} else {
			L.PushBool(true)
			return 1
		}
	}
}

func SetLuaFunctions(this *Lua) {
	stackPos := this.GetTop()
	defer this.SetTop(stackPos)
	this.NewTable()
	this.PushGoFunction(cmdAlias)
	this.SetField(-2, "alias")
	this.PushGoFunction(cmdSetEnv)
	this.SetField(-2, "setenv")
	this.PushGoFunction(cmdGetEnv)
	this.SetField(-2, "getenv")
	this.PushGoFunction(cmdExec)
	this.SetField(-2, "exec")
	this.PushGoFunction(cmdWrite)
	this.SetField(-2, "write")
	this.PushGoFunction(cmdAtoU)
	this.SetField(-2, "atou")
	this.PushGoFunction(cmdUtoA)
	this.SetField(-2, "utoa")
	this.PushGoFunction(cmdGetwd)
	this.SetField(-2, "getwd")
	this.PushGoFunction(cmdWhich)
	this.SetField(-2, "which")
	this.PushGoFunction(cmdEval)
	this.SetField(-2, "eval")
	this.PushGoFunction(cmdGlob)
	this.SetField(-2, "glob")
	this.PushGoFunction(cmdBindKey)
	this.SetField(-2, "bindkey")
	exeName, exeNameErr := dos.GetModuleFileName()
	if exeNameErr != nil {
		fmt.Fprintln(os.Stderr, exeNameErr)
	} else {
		this.PushString(exeName)
		this.SetField(-2, "exe")
	}
	this.SetGlobal("nyagos")

	// replace io.getenv
	this.GetGlobal("os")
	this.PushGoFunction(cmdGetEnv)
	this.SetField(-2, "getenv")

	var orgArgHook func([]string) []string
	orgArgHook = interpreter.SetArgsHook(func(args []string) []string {
		pos := this.GetTop()
		defer this.SetTop(pos)
		this.GetGlobal("nyagos")
		this.GetField(-1, "argsfilter")
		if !this.IsFunction(-1) {
			return orgArgHook(args)
		}
		this.NewTable()
		for i := 0; i < len(args); i++ {
			this.PushInteger(i)
			this.PushString(args[i])
			this.SetTable(-3)
		}
		if err := this.Call(1, 1); err != nil {
			fmt.Fprintf(os.Stderr, "%s\n", err)
			return orgArgHook(args)
		}
		if this.GetType(-1) != TTABLE {
			return orgArgHook(args)
		}
		newargs := []string{}
		for i := 0; true; i++ {
			this.PushInteger(i)
			this.GetTable(-2)
			if this.GetType(-1) == TNIL {
				break
			}
			arg1, arg1err := this.ToString(-1)
			if arg1err == nil {
				newargs = append(newargs, arg1)
			} else {
				fmt.Fprintln(os.Stderr, arg1err.Error())
			}
			this.Pop(1)
		}
		return orgArgHook(newargs)
	})
}
<commit_msg>import . "./lua" -> import "./lua"<commit_after>package main

import "fmt"
import "io"
import "os"
import "os/exec"
import "strings"

import "./lua"
import "./alias"
import "./conio/readline"
import "./dos"
import "./interpreter"
import "./mbcs"

import "github.com/shiena/ansicolor"

type LuaFunction struct {
	L            *lua.Lua
	registoryKey string
}

var LuaInstanceToCmd = map[uintptr]*interpreter.Interpreter{}

func (this LuaFunction) String() string {
	return "<<Lua-function>>"
}

func (this LuaFunction) Call(cmd *interpreter.Interpreter) (interpreter.NextT, error) {
	this.L.GetField(lua.Registory, this.registoryKey)
	this.L.NewTable()
	for i, arg1 := range cmd.Args {
		this.L.PushInteger(i)
		this.L.PushString(arg1)
		this.L.SetTable(-3)
	}
	LuaInstanceToCmd[this.L.Id()] = cmd
	err := this.L.Call(1, 0)
	return interpreter.CONTINUE, err
}

func cmdAlias(L *lua.Lua) int {
	name, nameErr := L.ToString(1)
	if nameErr != nil {
		L.PushNil()
		L.PushString(nameErr.Error())
		return 2
	}
	key := strings.ToLower(name)
	switch L.GetType(2) {
	case lua.TSTRING:
		value, err := L.ToString(2)
		if err == nil {
			alias.Table[key] = alias.New(value)
		} else {
			L.PushNil()
			L.PushString(err.Error())
			return 2
		}
	case lua.TFUNCTION:
		regkey := "nyagos.alias." + key
		L.SetField(lua.Registory, regkey)
		alias.Table[key] = LuaFunction{L, regkey}
	}
	L.PushBool(true)
	return 1
}

func cmdSetEnv(L *lua.Lua) int {
	name, nameErr := L.ToString(1)
	if nameErr != nil {
		L.PushNil()
		L.PushString(nameErr.Error())
		return 2
	}
	value, valueErr := L.ToString(2)
	if valueErr != nil {
		L.PushNil()
		L.PushString(valueErr.Error())
		return 2
	}
	os.Setenv(name, value)
	L.PushBool(true)
	return 1
}

func cmdGetEnv(L *lua.Lua) int {
	name, nameErr := L.ToString(1)
	if nameErr != nil {
		L.PushNil()
		return 1
	}
	value := os.Getenv(name)
	if len(value) > 0 {
		L.PushString(value)
	} else {
		L.PushNil()
	}
	return 1
}

func cmdExec(L *lua.Lua) int {
	statement, statementErr := L.ToString(1)
	if statementErr != nil {
		L.PushNil()
		L.PushString(statementErr.Error())
		return 2
	}
	_, err := interpreter.New().Interpret(statement)

	if err != nil {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
	L.PushBool(true)
	return 1
}

func cmdEval(L *lua.Lua) int {
	statement, statementErr := L.ToString(1)
	if statementErr != nil {
		L.PushNil()
		L.PushString(statementErr.Error())
		return 2
	}
	r, w, err := os.Pipe()
	if err != nil {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
	go func(statement string, w *os.File) {
		it := interpreter.New()
		it.Stdout = w
		it.Interpret(statement)
		w.Close()
	}(statement, w)

	var result = []byte{}
	for {
		buffer := make([]byte, 256)
		size, err := r.Read(buffer)
		if err != nil || size <= 0 {
			break
		}
		result = append(result, buffer[0:size]...)
	}
	r.Close()
	L.PushAnsiString(result)
	return 1
}

func cmdWrite(L *lua.Lua) int {
	var out io.Writer = os.Stdout
	cmd, cmdOk := LuaInstanceToCmd[L.Id()]
	if cmdOk && cmd != nil && cmd.Stdout != nil {
		out = cmd.Stdout
	}
	switch out.(type) {
	case *os.File:
		out = ansicolor.NewAnsiColorWriter(out)
	}

	n := L.GetTop()
	for i := 1; i <= n; i++ {
		str, err := L.ToString(i)
		if err != nil {
			L.PushNil()
			L.PushString(err.Error())
			return 2
		}
		if i > 1 {
			fmt.Fprint(out, "\t")
		}
		fmt.Fprint(out, str)
	}
	L.PushBool(true)
	return 1
}

func cmdGetwd(L *lua.Lua) int {
	wd, err := os.Getwd()
	if err == nil {
		L.PushString(wd)
		return 1
	} else {
		return 0
	}
}

func cmdWhich(L *lua.Lua) int {
	if L.GetType(-1) != lua.TSTRING {
		return 0
	}
	name, nameErr := L.ToString(-1)
	if nameErr != nil {
		L.PushNil()
		L.PushString(nameErr.Error())
		return 2
	}
	path, err := exec.LookPath(name)
	if err == nil {
		L.PushString(path)
		return 1
	} else {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
}

func cmdAtoU(L *lua.Lua) int {
	str, err := mbcs.AtoU(L.ToAnsiString(1))
	if err == nil {
		L.PushString(str)
		return 1
	} else {
		return 0
	}
}

func cmdUtoA(L *lua.Lua) int {
	utf8, utf8err := L.ToString(1)
	if utf8err != nil {
		L.PushNil()
		L.PushString(utf8err.Error())
		return 2
	}
	str, err := mbcs.UtoA(utf8)
	if err == nil {
		if len(str) >= 1 {
			L.PushAnsiString(str[:len(str)-1])
		} else {
			L.PushString("")
		}
		L.PushNil()
		return 2
	} else {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
}

func cmdGlob(L *lua.Lua) int {
	if !L.IsString(-1) {
		return 0
	}
	wildcard, wildcardErr := L.ToString(-1)
	if wildcardErr != nil {
		L.PushNil()
		L.PushString(wildcardErr.Error())
		return 2
	}
	list, err := dos.Glob(wildcard)
	if err != nil {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	} else {
		L.NewTable()
		for i := 0; i < len(list); i++ {
			L.PushInteger(i + 1)
			L.PushString(list[i])
			L.SetTable(-3)
		}
		return 1
	}
}

type KeyLuaFuncT struct {
	L            *lua.Lua
	registoryKey string
}

func (this *KeyLuaFuncT) Call(buffer *readline.Buffer) readline.Result {
	this.L.GetField(lua.Registory, this.registoryKey)
	if err := this.L.Call(0, 0); err != nil {
		fmt.Println(os.Stderr, err)
	}
	return readline.CONTINUE
}

func cmdBindKey(L *lua.Lua) int {
	key, keyErr := L.ToString(-2)
	if keyErr != nil {
		L.PushString(keyErr.Error())
		return 1
	}
	key = strings.Replace(strings.ToUpper(key), "-", "_", -1)
	switch L.GetType(-1) {
	case lua.TFUNCTION:
		regkey := "nyagos.bind." + key
		L.SetField(lua.Registory, regkey)
		if err := readline.BindKeyFunc(key, &KeyLuaFuncT{L, regkey}); err != nil {
			L.PushNil()
			L.PushString(err.Error())
			return 2
		} else {
			L.PushBool(true)
			return 1
		}
	default:
		val, valErr := L.ToString(-1)
		if valErr != nil {
			L.PushNil()
			L.PushString(valErr.Error())
			return 2
		}
		err := readline.BindKeySymbol(key, val)
		if err != nil {
			L.PushNil()
			L.PushString(err.Error())
			return 2
		} else {
			L.PushBool(true)
			return 1
		}
	}
}

func SetLuaFunctions(this *lua.Lua) {
	stackPos := this.GetTop()
	defer this.SetTop(stackPos)
	this.NewTable()
	this.PushGoFunction(cmdAlias)
	this.SetField(-2, "alias")
	this.PushGoFunction(cmdSetEnv)
	this.SetField(-2, "setenv")
	this.PushGoFunction(cmdGetEnv)
	this.SetField(-2, "getenv")
	this.PushGoFunction(cmdExec)
	this.SetField(-2, "exec")
	this.PushGoFunction(cmdWrite)
	this.SetField(-2, "write")
	this.PushGoFunction(cmdAtoU)
	this.SetField(-2, "atou")
	this.PushGoFunction(cmdUtoA)
	this.SetField(-2, "utoa")
	this.PushGoFunction(cmdGetwd)
	this.SetField(-2, "getwd")
	this.PushGoFunction(cmdWhich)
	this.SetField(-2, "which")
	this.PushGoFunction(cmdEval)
	this.SetField(-2, "eval")
	this.PushGoFunction(cmdGlob)
	this.SetField(-2, "glob")
	this.PushGoFunction(cmdBindKey)
	this.SetField(-2, "bindkey")
	exeName, exeNameErr := dos.GetModuleFileName()
	if exeNameErr != nil {
		fmt.Fprintln(os.Stderr, exeNameErr)
	} else {
		this.PushString(exeName)
		this.SetField(-2, "exe")
	}
	this.SetGlobal("nyagos")

	// replace io.getenv
	this.GetGlobal("os")
	this.PushGoFunction(cmdGetEnv)
	this.SetField(-2, "getenv")

	var orgArgHook func([]string) []string
	orgArgHook = interpreter.SetArgsHook(func(args []string) []string {
		pos := this.GetTop()
		defer this.SetTop(pos)
		this.GetGlobal("nyagos")
		this.GetField(-1, "argsfilter")
		if !this.IsFunction(-1) {
			return orgArgHook(args)
		}
		this.NewTable()
		for i := 0; i < len(args); i++ {
			this.PushInteger(i)
			this.PushString(args[i])
			this.SetTable(-3)
		}
		if err := this.Call(1, 1); err != nil {
			fmt.Fprintf(os.Stderr, "%s\n", err)
			return orgArgHook(args)
		}
		if this.GetType(-1) != lua.TTABLE {
			return orgArgHook(args)
		}
		newargs := []string{}
		for i := 0; true; i++ {
			this.PushInteger(i)
			this.GetTable(-2)
			if this.GetType(-1) == lua.TNIL {
				break
			}
			arg1, arg1err := this.ToString(-1)
			if arg1err == nil {
				newargs = append(newargs, arg1)
			} else {
				fmt.Fprintln(os.Stderr, arg1err.Error())
			}
			this.Pop(1)
		}
		return orgArgHook(newargs)
	})
}
<|endoftext|>


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

package unix

import (
	"syscall"
	"unsafe"
)

//go:cgo_import_dynamic libc_getentropy getentropy "/usr/lib/libSystem.B.dylib"

func libc_getentropy_trampoline()

// GetEntropy calls the macOS getentropy system call.
func GetEntropy(p []byte) error {
	_, _, errno := syscall_syscall(funcPC(libc_getentropy_trampoline),
		uintptr(unsafe.Pointer(&p[0])),
		uintptr(len(p)),
		0)
	if errno != 0 {
		return errno
	}
	return nil
}

//go:linkname syscall_syscall syscall.syscall
func syscall_syscall(fn, a1, a2, a3 uintptr) (r1, r2 uintptr, err syscall.Errno)

//go:linkname funcPC runtime.funcPC
func funcPC(f interface{}) uintptr
<commit_msg>internal/syscall/unix: use internal/abi.FuncPC for syscall wrapper<commit_after>// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package unix

import (
	"internal/abi"
	"syscall"
	"unsafe"
)

//go:cgo_import_dynamic libc_getentropy getentropy "/usr/lib/libSystem.B.dylib"

func libc_getentropy_trampoline()

// GetEntropy calls the macOS getentropy system call.
func GetEntropy(p []byte) error {
	_, _, errno := syscall_syscall(abi.FuncPCABI0(libc_getentropy_trampoline),
		uintptr(unsafe.Pointer(&p[0])),
		uintptr(len(p)),
		0)
	if errno != 0 {
		return errno
	}
	return nil
}

//go:linkname syscall_syscall syscall.syscall
func syscall_syscall(fn, a1, a2, a3 uintptr) (r1, r2 uintptr, err syscall.Errno)
<|endoftext|>


<commit_before>package lz4

/*
#cgo LDFLAGS: -llz4
#include "lz4.h"
#include "lz4hc.h"
*/
import "C"

import (
	"encoding/binary"
	"errors"
	"fmt"
	"hash"
	"io"
	"unsafe"

	"github.com/vova616/xxhash"
)

const (
	BestSpeed          = 3
	BestCompression    = 9
	DefaultCompression = -1
	lz4EOM             = uint32(0)
	lz4Magic           = uint32(0x184D2204)
	lz4BlockSizeID     = 7
	lz4BlockSize       = 1 << (8 + (2 * lz4BlockSizeID))
)

var lz4Header = []byte{
	0x4, 0x22, 0x4d, 0x18, 0x64, 0x70, 0xb9,
}

func blockSize(blockId uint32) uint32 {
	return (1 << (8 + (2 * blockId)))
}

type writer struct {
	level      int
	err        error
	compressor func(src []byte, dst []byte, maxSize uint32) (int, error)

	h hash.Hash32
	w io.Writer
}

// NewWriter creates a new Writer that satisfies writes by compressing data
// written to w.
func NewWriter(w io.Writer) io.WriteCloser {
	z, _ := NewWriterLevel(w, DefaultCompression)
	return z
}

// NewWriterLevel is like NewWriter but specifies the compression level instead
// of assuming DefaultCompression.
func NewWriterLevel(w io.Writer, level int) (io.WriteCloser, error) {
	if level < DefaultCompression || level > BestCompression {
		return nil, fmt.Errorf("lz4: invalid compression level: %d", level)
	}
	return &writer{
		level: level,
		w:     w,
		h:     xxhash.New(0),
	}, nil
}

func (z *writer) writeHeader() error {
	// Write magic number and header
	if _, err := z.w.Write(lz4Header); err != nil {
		return err
	}
	return nil
}

func (z *writer) write(v interface{}) error {
	return binary.Write(z.w, binary.LittleEndian, v)
}

// Write writes a compressed form of p to the underlying io.Writer.
func (z *writer) Write(p []byte) (int, error) {
	if z.err != nil {
		return 0, z.err
	}
	// Write headers
	if z.compressor == nil {
		if z.level == BestCompression {
			z.compressor = lz4CompressBest
		} else {
			z.compressor = lz4CompressSpeed
		}
		z.err = z.writeHeader()
		if z.err != nil {
			return 0, z.err
		}
	}

	compressed := make([]byte, lz4BlockSize)
	n, err := z.compressor(p, compressed, lz4BlockSize)
	if err != nil {
		z.err = err
		return 0, z.err
	}

	if n > 0 {
		z.err = z.write(uint32(n))
		if z.err != nil {
			return 0, z.err
		}
		// Write compressed block
		_, z.err = z.w.Write(compressed[0:n])
		if z.err != nil {
			return 0, z.err
		}
	} else {
		z.err = z.write(uint32(len(p)) | 0x80000000)
		if z.err != nil {
			return 0, z.err
		}
		// Write uncompressed block
		_, z.err = z.w.Write(p)
		if z.err != nil {
			return 0, z.err
		}
	}

	if len(p) > 0 {
		z.h.Write(p)
	}

	return len(p), nil
}

// Close closes the Writer. It does not close the underlying io.Writer.
func (z *writer) Close() error {
	z.err = z.write(uint32(0))
	if z.err != nil {
		return z.err
	}
	return z.write(z.h.Sum32())
}

func lz4CompressSpeed(src []byte, dst []byte, maxSize uint32) (int, error) {
	if len(src) == 0 {
		return 0, nil
	}
	n := C.LZ4_compress_limitedOutput((*C.char)(unsafe.Pointer(&src[0])), (*C.char)(unsafe.Pointer(&dst[0])), C.int(len(src)), C.int(maxSize))
	if n <= 0 {
		return 0, errors.New("lz4: data corruption")
	}
	return int(n), nil
}

func lz4CompressBest(src []byte, dst []byte, maxSize uint32) (int, error) {
	n := C.LZ4_compressHC_limitedOutput((*C.char)(unsafe.Pointer(&src[0])), (*C.char)(unsafe.Pointer(&dst[0])), C.int(len(src)), C.int(maxSize))
	if n <= 0 {
		return 0, errors.New("lz4: data corruption")
	}
	return int(n), nil
}

type reader struct {
	maxBlockSize        uint32
	contentChecksumFlag bool
	blockChecksumFlag   bool

	buf []byte
	r   io.Reader
	h   hash.Hash32
	err error
}

// NewReader creates a new Reader reading the given reader.
func NewReader(r io.Reader) (io.ReadCloser, error) {
	z := &reader{
		r: r,
		h: xxhash.New(0),
	}
	if err := z.readFrame(); err != nil {
		return nil, err
	}
	return z, nil
}

func (z *reader) readFrame() error {
	// Read and check magic
	var magic uint32
	if err := z.read(&magic); err != nil {
		return err
	}
	if magic != lz4Magic {
		return errors.New("lz4: invalid header")
	}
	br := newBitReader(io.LimitReader(z.r, 3))
	version, err := br.ReadBits(2)
	if err != nil {
		return err
	}
	if version != 1 {
		return errors.New("lz4: wrong version number")
	}
	independenceFlag, err := br.ReadBit()
	if err != nil || !independenceFlag {
		return err
	}
	blockChecksumFlag, err := br.ReadBit()
	if err != nil {
		return err
	}
	z.blockChecksumFlag = blockChecksumFlag
	contentSizeFlag, err := br.ReadBit()
	if err != nil || contentSizeFlag {
		return errors.New("lz4: does not support stream size")
	}
	contentChecksumFlag, err := br.ReadBit()
	if err != nil {
		return err
	}
	z.contentChecksumFlag = contentChecksumFlag
	if reserved, err := br.ReadBit(); err != nil || reserved {
		return errors.New("lz4: wrong value for reserved bits")
	}
	dictionaryFlag, err := br.ReadBit()
	if err != nil {
		return err
	}
	if dictionaryFlag {
		return errors.New("lz4: does not support dictionary")
	}
	if reserved, err := br.ReadBit(); err != nil || reserved {
		return errors.New("lz4: wrong value for reserved bits")
	}
	blockMaxSize, err := br.ReadBits(3)
	if err != nil {
		return err
	}
	if blockMaxSize < 4 {
		return errors.New("lz4: unsupported block size")
	}
	z.maxBlockSize = blockSize(blockMaxSize)
	if reserved, err := br.ReadBits(4); err != nil || reserved != 0 {
		return errors.New("lz4: wrong value for reserved bits")
	}
	sum := br.Sum32() >> 8 & 0xFF
	checksum, err := br.ReadBits(8)
	if err != nil {
		return err
	}
	if checksum != sum {
		return errors.New("lz4: stream descriptor error detected")
	}
	z.h.Reset()
	return nil
}

func (z *reader) nextBlock() {
	// Read block size
	var blockSize uint32
	z.err = z.read(&blockSize)
	if z.err != nil {
		return
	}

	uncompressedFlag := (blockSize >> 31) != 0
	blockSize &= 0x7FFFFFFF

	if blockSize == lz4EOM {
		z.err = io.EOF
		return
	}

	if blockSize > z.maxBlockSize {
		z.err = errors.New("lz4: invalid block size")
		return
	}

	// Read block data
	block := make([]byte, blockSize)
	_, z.err = io.ReadFull(z.r, block)
	if z.err != nil {
		return
	}

	if z.blockChecksumFlag {
		// Check block checksum
		var checksum uint32
		z.err = z.read(&checksum)
		if z.err != nil {
			return
		}
		if checksum != xxhash.Checksum32(block) {
			z.err = errors.New("lz4: invalid block checksum detected")
			return
		}
	}

	// Decompress
	data := make([]byte, z.maxBlockSize)
	if !uncompressedFlag {
		n, err := lz4Decompress(block, data, z.maxBlockSize)
		if err != nil {
			z.err = err
			return
		}
		data = data[0:n]
	} else {
		copy(data, block)
		data = data[0:blockSize]
	}

	if z.contentChecksumFlag {
		z.h.Write(data)
	}

	// Add block to our history
	z.buf = append(z.buf, data...)
}

func (z *reader) read(data interface{}) error {
	return binary.Read(z.r, binary.LittleEndian, data)
}

// Read reads a decompressed form of p from the underlying io.Reader.
func (z *reader) Read(p []byte) (int, error) {
	for {
		if len(z.buf) > 0 {
			n := copy(p, z.buf)
			z.buf = z.buf[n:]
			return n, nil
		}
		if z.err != nil {
			return 0, z.err
		}
		z.nextBlock()
	}
}

// Close closes the Reader. It does not close the underlying io.Reader.
func (z *reader) Close() error {
	if z.contentChecksumFlag {
		// Check content checksum
		var checksum uint32
		z.err = z.read(&checksum)
		if z.err != nil || checksum != z.h.Sum32() {
			z.err = errors.New("lz4: invalid content checksum detected")
			return z.err
		}
	}
	if z.err == io.EOF {
		return nil
	}
	return z.err
}

func lz4Decompress(src []byte, dst []byte, maxSize uint32) (int, error) {
	n := C.LZ4_decompress_safe((*C.char)(unsafe.Pointer(&src[0])),
		(*C.char)(unsafe.Pointer(&dst[0])), C.int(len(src)), C.int(maxSize))
	if n < 0 {
		return 0, errors.New("lz4: data corruption")
	}
	return int(n), nil
}
<commit_msg>rename blockId to blockID<commit_after>package lz4

/*
#cgo LDFLAGS: -llz4
#include "lz4.h"
#include "lz4hc.h"
*/
import "C"

import (
	"encoding/binary"
	"errors"
	"fmt"
	"hash"
	"io"
	"unsafe"

	"github.com/vova616/xxhash"
)

const (
	BestSpeed          = 3
	BestCompression    = 9
	DefaultCompression = -1
	lz4EOM             = uint32(0)
	lz4Magic           = uint32(0x184D2204)
	lz4BlockSizeID     = 7
	lz4BlockSize       = 1 << (8 + (2 * lz4BlockSizeID))
)

var lz4Header = []byte{
	0x4, 0x22, 0x4d, 0x18, 0x64, 0x70, 0xb9,
}

func blockSize(blockID uint32) uint32 {
	return (1 << (8 + (2 * blockID)))
}

type writer struct {
	level      int
	err        error
	compressor func(src []byte, dst []byte, maxSize uint32) (int, error)

	h hash.Hash32
	w io.Writer
}

// NewWriter creates a new Writer that satisfies writes by compressing data
// written to w.
func NewWriter(w io.Writer) io.WriteCloser {
	z, _ := NewWriterLevel(w, DefaultCompression)
	return z
}

// NewWriterLevel is like NewWriter but specifies the compression level instead
// of assuming DefaultCompression.
func NewWriterLevel(w io.Writer, level int) (io.WriteCloser, error) {
	if level < DefaultCompression || level > BestCompression {
		return nil, fmt.Errorf("lz4: invalid compression level: %d", level)
	}
	return &writer{
		level: level,
		w:     w,
		h:     xxhash.New(0),
	}, nil
}

func (z *writer) writeHeader() error {
	// Write magic number and header
	if _, err := z.w.Write(lz4Header); err != nil {
		return err
	}
	return nil
}

func (z *writer) write(v interface{}) error {
	return binary.Write(z.w, binary.LittleEndian, v)
}

// Write writes a compressed form of p to the underlying io.Writer.
func (z *writer) Write(p []byte) (int, error) {
	if z.err != nil {
		return 0, z.err
	}
	// Write headers
	if z.compressor == nil {
		if z.level == BestCompression {
			z.compressor = lz4CompressBest
		} else {
			z.compressor = lz4CompressSpeed
		}
		z.err = z.writeHeader()
		if z.err != nil {
			return 0, z.err
		}
	}

	compressed := make([]byte, lz4BlockSize)
	n, err := z.compressor(p, compressed, lz4BlockSize)
	if err != nil {
		z.err = err
		return 0, z.err
	}

	if n > 0 {
		z.err = z.write(uint32(n))
		if z.err != nil {
			return 0, z.err
		}
		// Write compressed block
		_, z.err = z.w.Write(compressed[0:n])
		if z.err != nil {
			return 0, z.err
		}
	} else {
		z.err = z.write(uint32(len(p)) | 0x80000000)
		if z.err != nil {
			return 0, z.err
		}
		// Write uncompressed block
		_, z.err = z.w.Write(p)
		if z.err != nil {
			return 0, z.err
		}
	}

	if len(p) > 0 {
		z.h.Write(p)
	}

	return len(p), nil
}

// Close closes the Writer. It does not close the underlying io.Writer.
func (z *writer) Close() error {
	z.err = z.write(uint32(0))
	if z.err != nil {
		return z.err
	}
	return z.write(z.h.Sum32())
}

func lz4CompressSpeed(src []byte, dst []byte, maxSize uint32) (int, error) {
	if len(src) == 0 {
		return 0, nil
	}
	n := C.LZ4_compress_limitedOutput((*C.char)(unsafe.Pointer(&src[0])), (*C.char)(unsafe.Pointer(&dst[0])), C.int(len(src)), C.int(maxSize))
	if n <= 0 {
		return 0, errors.New("lz4: data corruption")
	}
	return int(n), nil
}

func lz4CompressBest(src []byte, dst []byte, maxSize uint32) (int, error) {
	n := C.LZ4_compressHC_limitedOutput((*C.char)(unsafe.Pointer(&src[0])), (*C.char)(unsafe.Pointer(&dst[0])), C.int(len(src)), C.int(maxSize))
	if n <= 0 {
		return 0, errors.New("lz4: data corruption")
	}
	return int(n), nil
}

type reader struct {
	maxBlockSize        uint32
	contentChecksumFlag bool
	blockChecksumFlag   bool

	buf []byte
	r   io.Reader
	h   hash.Hash32
	err error
}

// NewReader creates a new Reader reading the given reader.
func NewReader(r io.Reader) (io.ReadCloser, error) {
	z := &reader{
		r: r,
		h: xxhash.New(0),
	}
	if err := z.readFrame(); err != nil {
		return nil, err
	}
	return z, nil
}

func (z *reader) readFrame() error {
	// Read and check magic
	var magic uint32
	if err := z.read(&magic); err != nil {
		return err
	}
	if magic != lz4Magic {
		return errors.New("lz4: invalid header")
	}
	br := newBitReader(io.LimitReader(z.r, 3))
	version, err := br.ReadBits(2)
	if err != nil {
		return err
	}
	if version != 1 {
		return errors.New("lz4: wrong version number")
	}
	independenceFlag, err := br.ReadBit()
	if err != nil || !independenceFlag {
		return err
	}
	blockChecksumFlag, err := br.ReadBit()
	if err != nil {
		return err
	}
	z.blockChecksumFlag = blockChecksumFlag
	contentSizeFlag, err := br.ReadBit()
	if err != nil || contentSizeFlag {
		return errors.New("lz4: does not support stream size")
	}
	contentChecksumFlag, err := br.ReadBit()
	if err != nil {
		return err
	}
	z.contentChecksumFlag = contentChecksumFlag
	if reserved, err := br.ReadBit(); err != nil || reserved {
		return errors.New("lz4: wrong value for reserved bits")
	}
	dictionaryFlag, err := br.ReadBit()
	if err != nil {
		return err
	}
	if dictionaryFlag {
		return errors.New("lz4: does not support dictionary")
	}
	if reserved, err := br.ReadBit(); err != nil || reserved {
		return errors.New("lz4: wrong value for reserved bits")
	}
	blockMaxSize, err := br.ReadBits(3)
	if err != nil {
		return err
	}
	if blockMaxSize < 4 {
		return errors.New("lz4: unsupported block size")
	}
	z.maxBlockSize = blockSize(blockMaxSize)
	if reserved, err := br.ReadBits(4); err != nil || reserved != 0 {
		return errors.New("lz4: wrong value for reserved bits")
	}
	sum := br.Sum32() >> 8 & 0xFF
	checksum, err := br.ReadBits(8)
	if err != nil {
		return err
	}
	if checksum != sum {
		return errors.New("lz4: stream descriptor error detected")
	}
	z.h.Reset()
	return nil
}

func (z *reader) nextBlock() {
	// Read block size
	var blockSize uint32
	z.err = z.read(&blockSize)
	if z.err != nil {
		return
	}

	uncompressedFlag := (blockSize >> 31) != 0
	blockSize &= 0x7FFFFFFF

	if blockSize == lz4EOM {
		z.err = io.EOF
		return
	}

	if blockSize > z.maxBlockSize {
		z.err = errors.New("lz4: invalid block size")
		return
	}

	// Read block data
	block := make([]byte, blockSize)
	_, z.err = io.ReadFull(z.r, block)
	if z.err != nil {
		return
	}

	if z.blockChecksumFlag {
		// Check block checksum
		var checksum uint32
		z.err = z.read(&checksum)
		if z.err != nil {
			return
		}
		if checksum != xxhash.Checksum32(block) {
			z.err = errors.New("lz4: invalid block checksum detected")
			return
		}
	}

	// Decompress
	data := make([]byte, z.maxBlockSize)
	if !uncompressedFlag {
		n, err := lz4Decompress(block, data, z.maxBlockSize)
		if err != nil {
			z.err = err
			return
		}
		data = data[0:n]
	} else {
		copy(data, block)
		data = data[0:blockSize]
	}

	if z.contentChecksumFlag {
		z.h.Write(data)
	}

	// Add block to our history
	z.buf = append(z.buf, data...)
}

func (z *reader) read(data interface{}) error {
	return binary.Read(z.r, binary.LittleEndian, data)
}

// Read reads a decompressed form of p from the underlying io.Reader.
func (z *reader) Read(p []byte) (int, error) {
	for {
		if len(z.buf) > 0 {
			n := copy(p, z.buf)
			z.buf = z.buf[n:]
			return n, nil
		}
		if z.err != nil {
			return 0, z.err
		}
		z.nextBlock()
	}
}

// Close closes the Reader. It does not close the underlying io.Reader.
func (z *reader) Close() error {
	if z.contentChecksumFlag {
		// Check content checksum
		var checksum uint32
		z.err = z.read(&checksum)
		if z.err != nil || checksum != z.h.Sum32() {
			z.err = errors.New("lz4: invalid content checksum detected")
			return z.err
		}
	}
	if z.err == io.EOF {
		return nil
	}
	return z.err
}

func lz4Decompress(src []byte, dst []byte, maxSize uint32) (int, error) {
	n := C.LZ4_decompress_safe((*C.char)(unsafe.Pointer(&src[0])),
		(*C.char)(unsafe.Pointer(&dst[0])), C.int(len(src)), C.int(maxSize))
	if n < 0 {
		return 0, errors.New("lz4: data corruption")
	}
	return int(n), nil
}
<|endoftext|>


<commit_before>package kontrol

import (
	"database/sql"
	"errors"
	"fmt"
	"net/url"
	"regexp"
	"strings"

	"github.com/hashicorp/go-version"
	_ "github.com/lib/pq"

	"github.com/koding/kite"
	kontrolprotocol "github.com/koding/kite/kontrol/protocol"
	"github.com/koding/kite/protocol"
	"github.com/koding/logging"
)

// Postgres holds Postgresql database related configuration
type PostgresConfig struct {
	Host     string
	Port     int
	Username string
	Password string
	DBName   string
}

type Postgres struct {
	DB  *sql.DB
	Log logging.Logger
}

func NewPostgres(conf *PostgresConfig, log kite.Logger) *Postgres {
	if conf.Port == 0 {
		conf.Port = 5432
	}

	if conf.Host == "" {
		conf.Host = "localhost"
	}

	if conf.DBName == "" {
		conf.DBName = "test"
	}

	connString := fmt.Sprintf(
		"host=%s port=%d dbname=%s sslmode=disable",
		conf.Host, conf.Port, conf.DBName,
	)

	if conf.Password != "" {
		connString += " password=" + conf.Password
	}

	if conf.Username != "" {
		connString += " user=" + conf.Username
	}

	db, err := sql.Open("postgres", connString)
	if err != nil {
		panic(err)
	}

	// add a limit so we don't hit a "too many open connections" errors. We
	// might change this in the future to tweak according to the machine and
	// usage behaviour
	db.SetMaxIdleConns(100)
	db.SetMaxOpenConns(100)

	// enable the ltree module which we are going to use, any error means it's
	// failed so there is no sense to continue, panic!
	enableTree := `CREATE EXTENSION IF NOT EXISTS ltree`
	if _, err := db.Exec(enableTree); err != nil {
		panic(err)
	}

	// create our initial kites table
	// * kite is going to be our ltree
	// * url is containing the kite's register url
	// * id is going to be kites' unique id (which also exists in the ltree
	//   path). We are adding it as a primary key so each kite with the full
	//   path can only exist once.
	// * created_at and updated_at are updated at creation and updating (like
	//  if the URL has changed)
	// Some notes:
	// * path label can only contain a sequence of alphanumeric characters
	//   and underscores. So for example a version string of "1.0.4" needs to
	//   be converted to "1_0_4" or uuid of 1111-2222-3333-4444 needs to be
	//   converted to 1111_2222_3333_4444.
	table := `CREATE TABLE IF NOT EXISTS kites (
		kite ltree NOT NULL,
		url text NOT NULL,
		id uuid PRIMARY KEY,
		created_at timestamptz NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC'),
		updated_at timestamptz NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC')
	);`

	if _, err := db.Exec(table); err != nil {
		panic(err)
	}

	// We enable index on the kite field. We don't return on errors because the
	// operator `IF NOT EXISTS` doesn't work for index creation, therefore we
	// assume the indexes might be already created.
	enableGistIndex := `CREATE INDEX kite_path_gist_idx ON kites USING GIST(kite)`
	enableBtreeIndex := `CREATE INDEX kite_path_btree_idx ON kites USING BTREE(kite)`

	if _, err := db.Exec(enableGistIndex); err != nil {
		log.Warning("postgres: enable gist index: ", err)
	}

	if _, err := db.Exec(enableBtreeIndex); err != nil {
		log.Warning("postgres: enable btree index: ", err)
	}

	return &Postgres{
		DB: db,
	}
}

func (p *Postgres) Get(query *protocol.KontrolQuery) (Kites, error) {
	// only let query with usernames, otherwise the whole tree will be fetched
	// which is not good for us
	if query.Username == "" {
		return nil, errors.New("username is not specified in query")
	}

	// NewVersion returns an error if it's a constraint, like: ">= 1.0, < 1.4"
	_, err := version.NewVersion(query.Version)
	if err != nil && query.Version != "" {
		return nil, errors.New("version constraint is not implemented")
	}

	path := ltreePath(query)

	// TODO: cache the statement or throttle it in a safe way, for example we
	// could cache it for every incoming 10 Get call
	stmt, err := p.DB.Prepare(`SELECT kite, url FROM kites WHERE kite <@ $1`)
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	rows, err := stmt.Query(path)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var kitePath string
	var url string

	kites := make(Kites, 0)

	for rows.Next() {
		err := rows.Scan(&kitePath, &url)
		if err != nil {
			return nil, err
		}

		kiteProt, err := kiteFromPath(kitePath)
		if err != nil {
			return nil, err
		}

		kites = append(kites, &protocol.KiteWithToken{
			Kite: *kiteProt,
			URL:  url,
		})

		fmt.Printf("kitePath %+v\n", kitePath)
		fmt.Printf("url %+v\n", url)
	}

	if err := rows.Err(); err != nil {
		return nil, err
	}

	kites.Shuffle()

	return kites, nil
}

func (p *Postgres) Add(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) error {
	// check that the incoming URL is valid to prevent malformed input
	_, err := url.Parse(value.URL)
	if err != nil {
		return err
	}

	_, err = p.DB.Exec("INSERT into kites(kite, url, id) VALUES($1, $2, $3)",
		ltreePath(kiteProt.Query()),
		value.URL,
		kiteProt.ID,
	)
	return err
}

func (p *Postgres) Update(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) error {
	// check that the incoming url is valid to prevent malformed input
	_, err := url.Parse(value.URL)
	if err != nil {
		return err
	}

	// TODO: also consider just using WHERE id = kiteProt.ID, see how it's
	// performs out
	_, err = p.DB.Exec(`UPDATE kites SET url = $1, updated_at = (now() at time zone 'utc')
	WHERE kite ~ $2`,
		value.URL, ltreePath(kiteProt.Query()))

	return err
}

func (p *Postgres) Delete(kite *protocol.Kite) error {
	return errors.New("DELETE is not implemented")
}

// ltreeLabel satisfies a invalid ltree definition of a label in path.
// According to the definition it is: "A label is a sequence of alphanumeric
// characters and underscores (for example, in C locale the characters
// A-Za-z0-9_ are allowed). Labels must be less than 256 bytes long."
//
// We could express one character with "[A-Za-z0-9_]", a word with
// "[A-Za-z0-9_]+". However we want to catch words that are not valid labels so
// we negate them with the "^" character, so it will be : "[^[A-Za-z0-9_]]+".
// Finally we cann use the POSIX character class: [:word:] which is:
// "Alphanumeric characters plus "_"", so the final regexp will be
// "[^[:word]]+"
var invalidLabelRe = regexp.MustCompile("[^[:word:]]+")

// ltreePath returns a query path to be used with the ltree module in postgress
// in the form of "username.environment.kitename.version.region.hostname.id"
func ltreePath(query *protocol.KontrolQuery) string {
	path := ""
	fields := query.Fields()

	// we stop for the first empty value
	for _, key := range keyOrder {
		v := fields[key]
		if v == "" {
			break
		}

		// replace anything that doesn't match the definition for a ltree path
		// label with a underscore, so the version "0.0.1" will be "0_0_1", or
		// uuid of "1111-2222-3333-4444" will be converted to
		// 1111_2222_3333_4444. Strings that satisfies the requirement are
		// untouched.
		v = invalidLabelRe.ReplaceAllLiteralString(v, "_")

		path = path + v + "."
	}

	// remove the latest dot which causes an invalid query
	path = strings.TrimSuffix(path, ".")
	return path
}

// kiteFromPath returns a Query from the given ltree path label
func kiteFromPath(path string) (*protocol.Kite, error) {
	fields := strings.Split(path, ".")

	if len(fields) != 7 {
		return nil, fmt.Errorf("invalid ltree path: %s", path)
	}

	// those labels were converted by us, therefore convert them back
	version := strings.Replace(fields[3], "_", ".", -1)
	id := strings.Replace(fields[6], "_", "-", -1)

	return &protocol.Kite{
		Username:    fields[0],
		Environment: fields[1],
		Name:        fields[2],
		Version:     version,
		Region:      fields[4],
		Hostname:    fields[5],
		ID:          id,
	}, nil

}
<commit_msg>kontrol/postgres: set back to a query<commit_after>package kontrol

import (
	"database/sql"
	"errors"
	"fmt"
	"net/url"
	"regexp"
	"strings"

	"github.com/hashicorp/go-version"
	_ "github.com/lib/pq"

	"github.com/koding/kite"
	kontrolprotocol "github.com/koding/kite/kontrol/protocol"
	"github.com/koding/kite/protocol"
	"github.com/koding/logging"
)

// Postgres holds Postgresql database related configuration
type PostgresConfig struct {
	Host     string
	Port     int
	Username string
	Password string
	DBName   string
}

type Postgres struct {
	DB  *sql.DB
	Log logging.Logger
}

func NewPostgres(conf *PostgresConfig, log kite.Logger) *Postgres {
	if conf.Port == 0 {
		conf.Port = 5432
	}

	if conf.Host == "" {
		conf.Host = "localhost"
	}

	if conf.DBName == "" {
		conf.DBName = "test"
	}

	connString := fmt.Sprintf(
		"host=%s port=%d dbname=%s sslmode=disable",
		conf.Host, conf.Port, conf.DBName,
	)

	if conf.Password != "" {
		connString += " password=" + conf.Password
	}

	if conf.Username != "" {
		connString += " user=" + conf.Username
	}

	db, err := sql.Open("postgres", connString)
	if err != nil {
		panic(err)
	}

	// add a limit so we don't hit a "too many open connections" errors. We
	// might change this in the future to tweak according to the machine and
	// usage behaviour
	db.SetMaxIdleConns(100)
	db.SetMaxOpenConns(100)

	// enable the ltree module which we are going to use, any error means it's
	// failed so there is no sense to continue, panic!
	enableTree := `CREATE EXTENSION IF NOT EXISTS ltree`
	if _, err := db.Exec(enableTree); err != nil {
		panic(err)
	}

	// create our initial kites table
	// * kite is going to be our ltree
	// * url is containing the kite's register url
	// * id is going to be kites' unique id (which also exists in the ltree
	//   path). We are adding it as a primary key so each kite with the full
	//   path can only exist once.
	// * created_at and updated_at are updated at creation and updating (like
	//  if the URL has changed)
	// Some notes:
	// * path label can only contain a sequence of alphanumeric characters
	//   and underscores. So for example a version string of "1.0.4" needs to
	//   be converted to "1_0_4" or uuid of 1111-2222-3333-4444 needs to be
	//   converted to 1111_2222_3333_4444.
	table := `CREATE TABLE IF NOT EXISTS kites (
		kite ltree NOT NULL,
		url text NOT NULL,
		id uuid PRIMARY KEY,
		created_at timestamptz NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC'),
		updated_at timestamptz NOT NULL DEFAULT (NOW() AT TIME ZONE 'UTC')
	);`

	if _, err := db.Exec(table); err != nil {
		panic(err)
	}

	// We enable index on the kite field. We don't return on errors because the
	// operator `IF NOT EXISTS` doesn't work for index creation, therefore we
	// assume the indexes might be already created.
	enableGistIndex := `CREATE INDEX kite_path_gist_idx ON kites USING GIST(kite)`
	enableBtreeIndex := `CREATE INDEX kite_path_btree_idx ON kites USING BTREE(kite)`

	if _, err := db.Exec(enableGistIndex); err != nil {
		log.Warning("postgres: enable gist index: ", err)
	}

	if _, err := db.Exec(enableBtreeIndex); err != nil {
		log.Warning("postgres: enable btree index: ", err)
	}

	return &Postgres{
		DB: db,
	}
}

func (p *Postgres) Get(query *protocol.KontrolQuery) (Kites, error) {
	// only let query with usernames, otherwise the whole tree will be fetched
	// which is not good for us
	if query.Username == "" {
		return nil, errors.New("username is not specified in query")
	}

	// NewVersion returns an error if it's a constraint, like: ">= 1.0, < 1.4"
	_, err := version.NewVersion(query.Version)
	if err != nil && query.Version != "" {
		return nil, errors.New("version constraint is not implemented")
	}

	path := ltreePath(query)

	rows, err := p.DB.Query(`SELECT kite, url FROM kites WHERE kite <@ $1`, path)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var kitePath string
	var url string

	kites := make(Kites, 0)

	for rows.Next() {
		err := rows.Scan(&kitePath, &url)
		if err != nil {
			return nil, err
		}

		kiteProt, err := kiteFromPath(kitePath)
		if err != nil {
			return nil, err
		}

		kites = append(kites, &protocol.KiteWithToken{
			Kite: *kiteProt,
			URL:  url,
		})

		fmt.Printf("kitePath %+v\n", kitePath)
		fmt.Printf("url %+v\n", url)
	}

	if err := rows.Err(); err != nil {
		return nil, err
	}

	kites.Shuffle()

	return kites, nil
}

func (p *Postgres) Add(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) error {
	// check that the incoming URL is valid to prevent malformed input
	_, err := url.Parse(value.URL)
	if err != nil {
		return err
	}

	_, err = p.DB.Exec("INSERT into kites(kite, url, id) VALUES($1, $2, $3)",
		ltreePath(kiteProt.Query()),
		value.URL,
		kiteProt.ID,
	)
	return err
}

func (p *Postgres) Update(kiteProt *protocol.Kite, value *kontrolprotocol.RegisterValue) error {
	// check that the incoming url is valid to prevent malformed input
	_, err := url.Parse(value.URL)
	if err != nil {
		return err
	}

	// TODO: also consider just using WHERE id = kiteProt.ID, see how it's
	// performs out
	_, err = p.DB.Exec(`UPDATE kites SET url = $1, updated_at = (now() at time zone 'utc')
	WHERE kite ~ $2`,
		value.URL, ltreePath(kiteProt.Query()))

	return err
}

func (p *Postgres) Delete(kite *protocol.Kite) error {
	return errors.New("DELETE is not implemented")
}

// ltreeLabel satisfies a invalid ltree definition of a label in path.
// According to the definition it is: "A label is a sequence of alphanumeric
// characters and underscores (for example, in C locale the characters
// A-Za-z0-9_ are allowed). Labels must be less than 256 bytes long."
//
// We could express one character with "[A-Za-z0-9_]", a word with
// "[A-Za-z0-9_]+". However we want to catch words that are not valid labels so
// we negate them with the "^" character, so it will be : "[^[A-Za-z0-9_]]+".
// Finally we cann use the POSIX character class: [:word:] which is:
// "Alphanumeric characters plus "_"", so the final regexp will be
// "[^[:word]]+"
var invalidLabelRe = regexp.MustCompile("[^[:word:]]+")

// ltreePath returns a query path to be used with the ltree module in postgress
// in the form of "username.environment.kitename.version.region.hostname.id"
func ltreePath(query *protocol.KontrolQuery) string {
	path := ""
	fields := query.Fields()

	// we stop for the first empty value
	for _, key := range keyOrder {
		v := fields[key]
		if v == "" {
			break
		}

		// replace anything that doesn't match the definition for a ltree path
		// label with a underscore, so the version "0.0.1" will be "0_0_1", or
		// uuid of "1111-2222-3333-4444" will be converted to
		// 1111_2222_3333_4444. Strings that satisfies the requirement are
		// untouched.
		v = invalidLabelRe.ReplaceAllLiteralString(v, "_")

		path = path + v + "."
	}

	// remove the latest dot which causes an invalid query
	path = strings.TrimSuffix(path, ".")
	return path
}

// kiteFromPath returns a Query from the given ltree path label
func kiteFromPath(path string) (*protocol.Kite, error) {
	fields := strings.Split(path, ".")

	if len(fields) != 7 {
		return nil, fmt.Errorf("invalid ltree path: %s", path)
	}

	// those labels were converted by us, therefore convert them back
	version := strings.Replace(fields[3], "_", ".", -1)
	id := strings.Replace(fields[6], "_", "-", -1)

	return &protocol.Kite{
		Username:    fields[0],
		Environment: fields[1],
		Name:        fields[2],
		Version:     version,
		Region:      fields[4],
		Hostname:    fields[5],
		ID:          id,
	}, nil

}
<|endoftext|>


<commit_before>package encoders

import (
	"errors"
	"strconv"

	"code.cloudfoundry.org/lager"

	"github.com/3d0c/gmf"
	"github.com/snickers/snickers/db"
	"github.com/snickers/snickers/types"
)

func addStream(job types.Job, codecName string, oc *gmf.FmtCtx, ist *gmf.Stream) (int, int, error) {
	var codecContext *gmf.CodecCtx
	var ost *gmf.Stream

	codec, err := gmf.FindEncoder(codecName)
	if err != nil {
		return 0, 0, err
	}

	if ost = oc.NewStream(codec); ost == nil {
		return 0, 0, errors.New("unable to create stream in output context")
	}
	defer gmf.Release(ost)

	if codecContext = gmf.NewCodecCtx(codec); codecContext == nil {
		return 0, 0, errors.New("unable to create codec context")
	}
	defer gmf.Release(codecContext)

	// https://ffmpeg.org/pipermail/ffmpeg-devel/2008-January/046900.html
	if oc.IsGlobalHeader() {
		codecContext.SetFlag(gmf.CODEC_FLAG_GLOBAL_HEADER)
	}

	if codec.IsExperimental() {
		codecContext.SetStrictCompliance(gmf.FF_COMPLIANCE_EXPERIMENTAL)
	}

	if codecContext.Type() == gmf.AVMEDIA_TYPE_AUDIO {
		err := setAudioCtxParams(codecContext, ist, job)
		if err != nil {
			return 0, 0, err
		}
	}

	if codecContext.Type() == gmf.AVMEDIA_TYPE_VIDEO {
		err := setVideoCtxParams(codecContext, ist, job)
		if err != nil {
			return 0, 0, err
		}
	}

	if err := codecContext.Open(nil); err != nil {
		return 0, 0, err
	}

	ost.SetCodecCtx(codecContext)

	return ist.Index(), ost.Index(), nil
}

// FFMPEGEncode function is responsible for encoding the file
func FFMPEGEncode(logger lager.Logger, dbInstance db.Storage, jobID string) error {
	log := logger.Session("ffmpeg-encode")
	log.Info("started", lager.Data{"job": jobID})
	defer log.Info("finished")

	gmf.LogSetLevel(gmf.AV_LOG_FATAL)
	job, _ := dbInstance.RetrieveJob(jobID)
	stMap := make(map[int]int, 0)
	var lastDelta int64

	inputCtx, err := gmf.NewInputCtx(job.LocalSource)
	if err != nil {
		log.Error("input-failed", err)
		return err
	}
	defer inputCtx.CloseInputAndRelease()

	outputCtx, err := gmf.NewOutputCtx(job.LocalDestination)
	if err != nil {
		log.Error("output-failed", err)
		return err
	}
	defer outputCtx.CloseOutputAndRelease()

	job.Status = types.JobEncoding
	job.Details = "0%"
	dbInstance.UpdateJob(job.ID, job)

	srcVideoStream, _ := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_VIDEO)

	videoCodec := getCodec(job)

	log.Info("add-stream-start", lager.Data{"code": videoCodec})
	i, o, err := addStream(job, videoCodec, outputCtx, srcVideoStream)
	if err != nil {
		log.Error("add-stream-failed", err)
		return err
	}
	log.Info("add-stream-finished")
	stMap[i] = o

	srcAudioStream, err := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_AUDIO)
	if err != nil {
		return err
	}

	audioCodec := "aac"
	if job.Preset.Audio.Codec != "aac" {
		audioCodec = job.Preset.Audio.Codec
	}

	i, o, err = addStream(job, audioCodec, outputCtx, srcAudioStream)
	if err != nil {
		return err
	}
	stMap[i] = o

	if err := outputCtx.WriteHeader(); err != nil {
		return err
	}
	totalFrames := float64(srcVideoStream.NbFrames() + srcAudioStream.NbFrames())
	framesCount := float64(0)

	for packet := range inputCtx.GetNewPackets() {
		ist, err := inputCtx.GetStream(packet.StreamIndex())
		if err != nil {
			return err
		}
		ost, err := outputCtx.GetStream(stMap[ist.Index()])
		if err != nil {
			return err
		}

		for frame := range packet.Frames(ist.CodecCtx()) {
			if ost.IsAudio() {
				fsTb := gmf.AVR{Num: 1, Den: ist.CodecCtx().SampleRate()}
				outTb := gmf.AVR{Num: 1, Den: ist.CodecCtx().SampleRate()}

				frame.SetPts(packet.Pts())

				pts := gmf.RescaleDelta(ist.TimeBase(), frame.Pts(), fsTb.AVRational(), frame.NbSamples(), &lastDelta, outTb.AVRational())

				frame.
					SetNbSamples(ost.CodecCtx().FrameSize()).
					SetFormat(ost.CodecCtx().SampleFmt()).
					SetChannelLayout(ost.CodecCtx().ChannelLayout()).
					SetPts(pts)
			} else {
				frame.SetPts(ost.Pts)
			}

			if p, ready, _ := frame.EncodeNewPacket(ost.CodecCtx()); ready {
				if p.Pts() != gmf.AV_NOPTS_VALUE {
					p.SetPts(gmf.RescaleQ(p.Pts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				if p.Dts() != gmf.AV_NOPTS_VALUE {
					p.SetDts(gmf.RescaleQ(p.Dts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				p.SetStreamIndex(ost.Index())

				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			}

			ost.Pts++
			framesCount++
			percentage := string(strconv.FormatInt(int64(framesCount/totalFrames*100), 10) + "%")
			if percentage != job.Details {
				job.Details = percentage
				dbInstance.UpdateJob(job.ID, job)
			}
		}
		gmf.Release(packet)
	}

	for i := 0; i < outputCtx.StreamsCnt(); i++ {
		ist, err := inputCtx.GetStream(0)
		if err != nil {
			return err
		}
		ost, err := outputCtx.GetStream(stMap[ist.Index()])
		if err != nil {
			return err
		}

		frame := gmf.NewFrame()

		for {
			if p, ready, _ := frame.FlushNewPacket(ost.CodecCtx()); ready {
				if p.Pts() != gmf.AV_NOPTS_VALUE {
					p.SetPts(gmf.RescaleQ(p.Pts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				if p.Dts() != gmf.AV_NOPTS_VALUE {
					p.SetDts(gmf.RescaleQ(p.Dts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				p.SetStreamIndex(ost.Index())

				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			} else {
				gmf.Release(p)
				break
			}

			ost.Pts++
		}

		gmf.Release(frame)
	}
	if job.Details != "100%" {
		job.Details = "100%"
		dbInstance.UpdateJob(job.ID, job)
	}

	return nil
}

func getProfile(job types.Job) int {
	profiles := map[string]int{
		"baseline": gmf.FF_PROFILE_H264_BASELINE,
		"main":     gmf.FF_PROFILE_H264_MAIN,
		"high":     gmf.FF_PROFILE_H264_HIGH,
	}

	if job.Preset.Video.Profile != "" {
		return profiles[job.Preset.Video.Profile]
	}
	return gmf.FF_PROFILE_H264_MAIN
}

func getCodec(job types.Job) string {
	codecs := map[string]string{
		"h264":   "libx264",
		"vp8":    "libvpx",
		"vp9":    "libvpx-vp9",
		"theora": "libtheora",
	}

	if job.Preset.Video.Codec != "" {
		return codecs[job.Preset.Video.Codec]
	}
	return "libx264"
}

func getResolution(job types.Job, inputWidth int, inputHeight int) (int, int) {
	var width, height int
	if job.Preset.Video.Width == "" && job.Preset.Video.Height == "" {
		return inputWidth, inputHeight
	} else if job.Preset.Video.Width == "" {
		height, _ = strconv.Atoi(job.Preset.Video.Height)
		width = (inputWidth * height) / inputHeight
	} else if job.Preset.Video.Height == "" {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height = (inputHeight * width) / inputWidth
	} else {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height, _ = strconv.Atoi(job.Preset.Video.Height)
	}
	return width, height
}

func setAudiooCtxParams(codecContext *gmf.CodecCtx, ist *gmf.Stream, job types.Job) error {
	bitrate, err := strconv.Atoi(job.Preset.Audio.Bitrate)
	if err != nil {
		return err
	}

	codecContext.SetBitRate(bitrate)
	codecContext.SetSampleFmt(ist.CodecCtx().SampleFmt())
	codecContext.SetSampleRate(ist.CodecCtx().SampleRate())
	codecContext.SetChannels(ist.CodecCtx().Channels())
	codecContext.SelectChannelLayout()
	codecContext.SelectSampleRate()
	return nil
}

func setVideoCtxParams(codecContext *gmf.CodecCtx, ist *gmf.Stream, job types.Job) error {
	codecContext.SetTimeBase(gmf.AVR{Num: 1, Den: 25}) // what is this

	if job.Preset.Video.Codec == "h264" {
		profile := getProfile(job)
		codecContext.SetProfile(profile)
	}

	gop, err := strconv.Atoi(job.Preset.Video.GopSize)
	if err != nil {
		return err
	}

	width, height := getResolution(job, ist.CodecCtx().Width(), ist.CodecCtx().Height())

	bitrate, err := strconv.Atoi(job.Preset.Video.Bitrate)
	if err != nil {
		return err
	}

	codecContext.SetDimension(width, height)
	codecContext.SetGopSize(gop)
	codecContext.SetBitRate(bitrate)
	codecContext.SetPixFmt(ist.CodecCtx().PixFmt())

	return nil
}
<commit_msg>encoders/ffmpeg: fix tests<commit_after>package encoders

import (
	"errors"
	"strconv"

	"code.cloudfoundry.org/lager"

	"github.com/3d0c/gmf"
	"github.com/snickers/snickers/db"
	"github.com/snickers/snickers/types"
)

func addStream(job types.Job, codecName string, oc *gmf.FmtCtx, ist *gmf.Stream) (int, int, error) {
	var codecContext *gmf.CodecCtx
	var ost *gmf.Stream

	codec, err := gmf.FindEncoder(codecName)
	if err != nil {
		return 0, 0, err
	}

	if ost = oc.NewStream(codec); ost == nil {
		return 0, 0, errors.New("unable to create stream in output context")
	}
	defer gmf.Release(ost)

	if codecContext = gmf.NewCodecCtx(codec); codecContext == nil {
		return 0, 0, errors.New("unable to create codec context")
	}
	defer gmf.Release(codecContext)

	// https://ffmpeg.org/pipermail/ffmpeg-devel/2008-January/046900.html
	if oc.IsGlobalHeader() {
		codecContext.SetFlag(gmf.CODEC_FLAG_GLOBAL_HEADER)
	}

	if codec.IsExperimental() {
		codecContext.SetStrictCompliance(gmf.FF_COMPLIANCE_EXPERIMENTAL)
	}

	if codecContext.Type() == gmf.AVMEDIA_TYPE_AUDIO {
		err := setAudioCtxParams(codecContext, ist, job)
		if err != nil {
			return 0, 0, err
		}
	}

	if codecContext.Type() == gmf.AVMEDIA_TYPE_VIDEO {
		err := setVideoCtxParams(codecContext, ist, job)
		if err != nil {
			return 0, 0, err
		}
	}

	if err := codecContext.Open(nil); err != nil {
		return 0, 0, err
	}

	ost.SetCodecCtx(codecContext)

	return ist.Index(), ost.Index(), nil
}

// FFMPEGEncode function is responsible for encoding the file
func FFMPEGEncode(logger lager.Logger, dbInstance db.Storage, jobID string) error {
	log := logger.Session("ffmpeg-encode")
	log.Info("started", lager.Data{"job": jobID})
	defer log.Info("finished")

	gmf.LogSetLevel(gmf.AV_LOG_FATAL)
	job, _ := dbInstance.RetrieveJob(jobID)
	stMap := make(map[int]int, 0)
	var lastDelta int64

	inputCtx, err := gmf.NewInputCtx(job.LocalSource)
	if err != nil {
		log.Error("input-failed", err)
		return err
	}
	defer inputCtx.CloseInputAndRelease()

	outputCtx, err := gmf.NewOutputCtx(job.LocalDestination)
	if err != nil {
		log.Error("output-failed", err)
		return err
	}
	defer outputCtx.CloseOutputAndRelease()

	job.Status = types.JobEncoding
	job.Details = "0%"
	dbInstance.UpdateJob(job.ID, job)

	srcVideoStream, _ := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_VIDEO)

	videoCodec := getCodec(job)

	log.Info("add-stream-start", lager.Data{"code": videoCodec})
	i, o, err := addStream(job, videoCodec, outputCtx, srcVideoStream)
	if err != nil {
		log.Error("add-stream-failed", err)
		return err
	}
	log.Info("add-stream-finished")
	stMap[i] = o

	srcAudioStream, err := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_AUDIO)
	if err != nil {
		return err
	}

	audioCodec := "aac"
	if job.Preset.Audio.Codec != "aac" {
		audioCodec = job.Preset.Audio.Codec
	}

	i, o, err = addStream(job, audioCodec, outputCtx, srcAudioStream)
	if err != nil {
		return err
	}
	stMap[i] = o

	if err := outputCtx.WriteHeader(); err != nil {
		return err
	}
	totalFrames := float64(srcVideoStream.NbFrames() + srcAudioStream.NbFrames())
	framesCount := float64(0)

	for packet := range inputCtx.GetNewPackets() {
		ist, err := inputCtx.GetStream(packet.StreamIndex())
		if err != nil {
			return err
		}
		ost, err := outputCtx.GetStream(stMap[ist.Index()])
		if err != nil {
			return err
		}

		for frame := range packet.Frames(ist.CodecCtx()) {
			if ost.IsAudio() {
				fsTb := gmf.AVR{Num: 1, Den: ist.CodecCtx().SampleRate()}
				outTb := gmf.AVR{Num: 1, Den: ist.CodecCtx().SampleRate()}

				frame.SetPts(packet.Pts())

				pts := gmf.RescaleDelta(ist.TimeBase(), frame.Pts(), fsTb.AVRational(), frame.NbSamples(), &lastDelta, outTb.AVRational())

				frame.
					SetNbSamples(ost.CodecCtx().FrameSize()).
					SetFormat(ost.CodecCtx().SampleFmt()).
					SetChannelLayout(ost.CodecCtx().ChannelLayout()).
					SetPts(pts)
			} else {
				frame.SetPts(ost.Pts)
			}

			if p, ready, _ := frame.EncodeNewPacket(ost.CodecCtx()); ready {
				if p.Pts() != gmf.AV_NOPTS_VALUE {
					p.SetPts(gmf.RescaleQ(p.Pts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				if p.Dts() != gmf.AV_NOPTS_VALUE {
					p.SetDts(gmf.RescaleQ(p.Dts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				p.SetStreamIndex(ost.Index())

				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			}

			ost.Pts++
			framesCount++
			percentage := string(strconv.FormatInt(int64(framesCount/totalFrames*100), 10) + "%")
			if percentage != job.Details {
				job.Details = percentage
				dbInstance.UpdateJob(job.ID, job)
			}
		}
		gmf.Release(packet)
	}

	for i := 0; i < outputCtx.StreamsCnt(); i++ {
		ist, err := inputCtx.GetStream(0)
		if err != nil {
			return err
		}
		ost, err := outputCtx.GetStream(stMap[ist.Index()])
		if err != nil {
			return err
		}

		frame := gmf.NewFrame()

		for {
			if p, ready, _ := frame.FlushNewPacket(ost.CodecCtx()); ready {
				if p.Pts() != gmf.AV_NOPTS_VALUE {
					p.SetPts(gmf.RescaleQ(p.Pts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				if p.Dts() != gmf.AV_NOPTS_VALUE {
					p.SetDts(gmf.RescaleQ(p.Dts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				p.SetStreamIndex(ost.Index())

				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			} else {
				gmf.Release(p)
				break
			}

			ost.Pts++
		}

		gmf.Release(frame)
	}
	if job.Details != "100%" {
		job.Details = "100%"
		dbInstance.UpdateJob(job.ID, job)
	}

	return nil
}

func getProfile(job types.Job) int {
	profiles := map[string]int{
		"baseline": gmf.FF_PROFILE_H264_BASELINE,
		"main":     gmf.FF_PROFILE_H264_MAIN,
		"high":     gmf.FF_PROFILE_H264_HIGH,
	}

	if job.Preset.Video.Profile != "" {
		return profiles[job.Preset.Video.Profile]
	}
	return gmf.FF_PROFILE_H264_MAIN
}

func getCodec(job types.Job) string {
	codecs := map[string]string{
		"h264":   "libx264",
		"vp8":    "libvpx",
		"vp9":    "libvpx-vp9",
		"theora": "libtheora",
	}

	if job.Preset.Video.Codec != "" {
		return codecs[job.Preset.Video.Codec]
	}
	return "libx264"
}

func GetResolution(job types.Job, inputWidth int, inputHeight int) (int, int) {
	var width, height int
	if job.Preset.Video.Width == "" && job.Preset.Video.Height == "" {
		return inputWidth, inputHeight
	} else if job.Preset.Video.Width == "" {
		height, _ = strconv.Atoi(job.Preset.Video.Height)
		width = (inputWidth * height) / inputHeight
	} else if job.Preset.Video.Height == "" {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height = (inputHeight * width) / inputWidth
	} else {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height, _ = strconv.Atoi(job.Preset.Video.Height)
	}
	return width, height
}

func setAudioCtxParams(codecContext *gmf.CodecCtx, ist *gmf.Stream, job types.Job) error {
	bitrate, err := strconv.Atoi(job.Preset.Audio.Bitrate)
	if err != nil {
		return err
	}

	codecContext.SetBitRate(bitrate)
	codecContext.SetSampleFmt(ist.CodecCtx().SampleFmt())
	codecContext.SetSampleRate(ist.CodecCtx().SampleRate())
	codecContext.SetChannels(ist.CodecCtx().Channels())
	codecContext.SelectChannelLayout()
	codecContext.SelectSampleRate()
	return nil
}

func setVideoCtxParams(codecContext *gmf.CodecCtx, ist *gmf.Stream, job types.Job) error {
	codecContext.SetTimeBase(gmf.AVR{Num: 1, Den: 25}) // what is this

	if job.Preset.Video.Codec == "h264" {
		profile := getProfile(job)
		codecContext.SetProfile(profile)
	}

	gop, err := strconv.Atoi(job.Preset.Video.GopSize)
	if err != nil {
		return err
	}

	width, height := GetResolution(job, ist.CodecCtx().Width(), ist.CodecCtx().Height())

	bitrate, err := strconv.Atoi(job.Preset.Video.Bitrate)
	if err != nil {
		return err
	}

	codecContext.SetDimension(width, height)
	codecContext.SetGopSize(gop)
	codecContext.SetBitRate(bitrate)
	codecContext.SetPixFmt(ist.CodecCtx().PixFmt())

	return nil
}
<|endoftext|>