<commit_before>/* Goiardi configuration. */

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

// Package config parses command line flags and config files, and defines
// options used elsewhere in goiardi.
package config

import (
	"github.com/jessevdk/go-flags"
	"github.com/BurntSushi/toml"
	"os"
	"log"
	"fmt"
	"time"
	"path"
)

/* Master struct for configuration. */
type Conf struct {
	Ipaddress string
	Port int
	Hostname string
	ConfFile string `toml:"conf-file"`
	IndexFile string `toml:"index-file"`
	DataStoreFile string `toml:"data-file"`
	DebugLevel int `toml:"debug-level"`
	FreezeInterval int `toml:"freeze-interval"`
	FreezeData bool `toml:"freeze-data"`
	LogFile string `toml:"log-file"`
	UseAuth bool `toml:"use-auth"`
	TimeSlew string `toml:"time-slew"`
	TimeSlewDur time.Duration
	ConfRoot string `toml:"conf-root"`
	UseSSL bool `toml:"use-ssl"`
	SslCert string `toml:"ssl-cert"`
	SslKey string `toml:"ssl-key"`
	HttpsUrls bool `toml:"https-urls"`
	DisableWebUI bool `toml:"disable-webui"`
	UseMySQL bool `toml:"use-mysql"`
	MySQL MySQLdb `toml:"mysql"`
	LocalFstoreDir string `toml:"local-filestore-dir"`
}

// MySQL connection options
type MySQLdb struct {
	Username string
	Password string
	Protocol string
	Address string
	Port string
	Dbname string
	ExtraParams map[string]string `toml:"extra_params"`
}

/* Struct for command line options. */
type Options struct {
	Version bool `short:"v" long:"version" description:"Print version info."`
	Verbose []bool `short:"V" long:"verbose" description:"Show verbose debug information. (not implemented)"`
	ConfFile string `short:"c" long:"config" description:"Specify a config file to use."`
	Ipaddress string `short:"I" long:"ipaddress" description:"Listen on a specific IP address."`
	Hostname string `short:"H" long:"hostname" description:"Hostname to use for this server. Defaults to hostname reported by the kernel."`
	Port int `short:"P" long:"port" description:"Port to listen on. If port is set to 443, SSL will be activated. (default: 4545)"`
	IndexFile string `short:"i" long:"index-file" description:"File to save search index data to."`
	DataStoreFile string `short:"D" long:"data-file" description:"File to save data store data to."`
	FreezeInterval int `short:"F" long:"freeze-interval" description:"Interval in seconds to freeze in-memory data structures to disk (requires -i/--index-file and -D/--data-file options to be set). (Default 300 seconds/5 minutes.)"`
	LogFile string `short:"L" long:"log-file" description:"Log to file X"`
	TimeSlew string `long:"time-slew" description:"Time difference allowed between the server's clock at the time in the X-OPS-TIMESTAMP header. Formatted like 5m, 150s, etc. Defaults to 15m."`
	ConfRoot string `long:"conf-root" description:"Root directory for configs and certificates. Default: the directory the config file is in, or the current directory if no config file is set."`
	UseAuth bool `short:"A" long:"use-auth" description:"Use authentication. Default: false."`
	UseSSL bool `long:"use-ssl" description:"Use SSL for connections. If --port is set to 433, this will automatically be turned on. If it is set to 80, it will automatically be turned off. Default: off. Requires --ssl-cert and --ssl-key."`
	SslCert string `long:"ssl-cert" description:"SSL certificate file. If a relative path, will be set relative to --conf-root."`
	SslKey string `long:"ssl-key" description:"SSL key file. If a relative path, will be set relative to --conf-root."`
	HttpsUrls bool `long:"https-urls" description:"Use 'https://' in URLs to server resources if goiardi is not using SSL for its connections. Useful when goiardi is sitting behind a reverse proxy that uses SSL, but is communicating with the proxy over HTTP."`
	DisableWebUI bool `long:"disable-webui" description:"If enabled, disables connections and logins to goiardi over the webui interface."`
	UseMySQL bool `long:"use-mysql" description:"Use a MySQL database for data storage. Configure database options in the config file."`
	LocalFstoreDir string `long:"local-filestore-dir" description:"Directory to save uploaded files in. Optional when running in in-memory mode, *mandatory* for SQL mode."`
}

// The goiardi version.
const Version = "0.4.2"
// The chef version we're at least aiming for, even if it's not complete yet.
const ChefVersion = "11.0.8"

/* The general plan is to read the command-line options, then parse the config
 * file, fill in the config struct with those values, then apply the
 * command-line options to the config struct. We read the cli options first so
 * we know to look for a different config file if needed, but otherwise the
 * command line options override what's in the config file. */

func InitConfig() *Conf { return &Conf{ } }

// Conf struct with the options specified on the command line or in the config
// file.
var Config = InitConfig()

// Read and apply arguments from the command line.
func ParseConfigOptions() error {
	var opts = &Options{ }
	_, err := flags.Parse(opts)

	if err != nil {
		if err.(*flags.Error).Type == flags.ErrHelp {
			os.Exit(0)
		} else {
			log.Println(err)
			os.Exit(1)
		}
	}

	if opts.Version {
		fmt.Printf("goiardi version %s (aiming for compatibility with Chef Server version %s).\n", Version, ChefVersion)
		os.Exit(0)
	}

	/* Load the config file. Command-line options have precedence over
	 * config file options. */
	if opts.ConfFile != "" {
		if _, err := toml.DecodeFile(opts.ConfFile, Config); err != nil {
			panic(err)
			os.Exit(1)
		}
		Config.ConfFile = opts.ConfFile
		Config.FreezeData = false
	}

	if opts.Hostname != "" {
		Config.Hostname = opts.Hostname
	} else {
		if Config.Hostname == "" {
			Config.Hostname, err = os.Hostname()
			if err != nil {
				log.Println(err)
				Config.Hostname = "localhost"
			}
		}
	}

	if opts.DataStoreFile != "" {
		Config.DataStoreFile = opts.DataStoreFile
	}

	if opts.IndexFile != "" {
		Config.IndexFile = opts.IndexFile
	}

	if !((Config.DataStoreFile == "" && Config.IndexFile == "") || (Config.DataStoreFile != "" && Config.IndexFile != "")) {
		err := fmt.Errorf("-i and -D must either both be specified, or not specified.")
		panic(err)
		os.Exit(1)
	}

	if Config.IndexFile != "" && Config.DataStoreFile != "" {
		Config.FreezeData = true
	}

	if opts.LogFile != "" {
		Config.LogFile = opts.LogFile
	}
	if Config.LogFile != "" {
		lfp, lerr := os.Create(Config.LogFile)
		if lerr != nil {
			log.Println(err)
			os.Exit(1)
		}
		log.SetOutput(lfp)
	}

	/* Database options */
	if opts.UseMySQL {
		Config.UseMySQL = opts.UseMySQL
	}
	// Don't bother setting a default mysql port if mysql isn't used
	if Config.UseMySQL {
		if Config.MySQL.Port == "" {
			Config.MySQL.Port = "3306"
		}
	}

	if opts.LocalFstoreDir != "" {
		Config.LocalFstoreDir = opts.LocalFstoreDir
	}
	if Config.LocalFstoreDir == "" && Config.UseMySQL {
		err := fmt.Errorf("local-filestore-dir must be set when running goiardi in SQL mode")
		log.Println(err)
		os.Exit(1)
	}

	// TODO: once db support is more mature, the dance with freezing data
	// will become more complicated and need to be changed.

	if !Config.FreezeData && (opts.FreezeInterval != 0 || Config.FreezeInterval != 0) {
		log.Printf("FYI, setting the freeze data interval's not especially useful without setting the index and data files.")
	}
	if opts.FreezeInterval != 0 {
		Config.FreezeInterval = opts.FreezeInterval
	}
	if Config.FreezeInterval == 0 {
		Config.FreezeInterval = 300
	}

	/* Root directory for certs and the like */
	if opts.ConfRoot != "" {
		Config.ConfRoot = opts.ConfRoot
	}

	if Config.ConfRoot == "" {
		if Config.ConfFile != "" {
			Config.ConfRoot = path.Dir(Config.ConfFile)
		} else {
			Config.ConfRoot = "."
		}
	}

	Config.Ipaddress = opts.Ipaddress
	if opts.Port != 0 {
		Config.Port = opts.Port
	}
	if Config.Port == 0 {
		Config.Port = 4545
	}

	if opts.UseSSL {
		Config.UseSSL = opts.UseSSL
	}
	if opts.SslCert != "" {
		Config.SslCert = opts.SslCert
	}
	if opts.SslKey != "" {
		Config.SslKey = opts.SslKey
	}
	if opts.HttpsUrls {
		Config.HttpsUrls = opts.HttpsUrls
	}
	// SSL setup
	if Config.Port == 80 {
		Config.UseSSL = false
	} else if Config.Port == 443 {
		Config.UseSSL = true
	}
	if Config.UseSSL {
		if Config.SslCert == "" || Config.SslKey == "" {
			log.Println("SSL mode requires specifying both a certificate and a key file.")
			os.Exit(1)
		}
		/* If the SSL cert and key are not absolute files, join them
		 * with the conf root */
		if !path.IsAbs(Config.SslCert) {
			Config.SslCert = path.Join(Config.ConfRoot, Config.SslCert)
		}
		if !path.IsAbs(Config.SslKey) {
			Config.SslKey = path.Join(Config.ConfRoot, Config.SslKey)
		}
	}

	Config.DebugLevel = len(opts.Verbose)

	if opts.TimeSlew != "" {
		Config.TimeSlew = opts.TimeSlew
	}
	if Config.TimeSlew != "" {
		d, derr := time.ParseDuration(Config.TimeSlew)
		if derr != nil {
			log.Println("Error parsing time-slew:", derr)
			os.Exit(1)
		}
		Config.TimeSlewDur = d
	} else {
		Config.TimeSlewDur, _ = time.ParseDuration("15m")
	}

	if opts.UseAuth {
		Config.UseAuth = opts.UseAuth
	}

	if opts.DisableWebUI {
		Config.DisableWebUI = opts.DisableWebUI
	}

	return nil
}

// The address and port goiardi is configured to listen on.
func ListenAddr() string {
	listen_addr := fmt.Sprintf("%s:%d", Config.Ipaddress, Config.Port)
	return listen_addr
}

// The hostname and port goiardi is configured to use.
func ServerHostname() string {
	var portStr string
	if !(Config.Port == 80 || Config.Port == 443) {
		portStr = fmt.Sprintf(":%d", Config.Port)
	}
	hostname := fmt.Sprintf("%s%s", Config.Hostname, portStr)
	return hostname
}

// The base URL
func ServerBaseURL() string {
	/* TODO: allow configuring using http vs. https */
	var urlScheme string
	if Config.UseSSL || Config.HttpsUrls {
		urlScheme = "https"
	} else {
		urlScheme = "http"
	}
	url := fmt.Sprintf("%s://%s", urlScheme, ServerHostname())
	return url
}
<commit_msg>Prerelease version<commit_after>/* Goiardi configuration. */

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

// Package config parses command line flags and config files, and defines
// options used elsewhere in goiardi.
package config

import (
	"github.com/jessevdk/go-flags"
	"github.com/BurntSushi/toml"
	"os"
	"log"
	"fmt"
	"time"
	"path"
)

/* Master struct for configuration. */
type Conf struct {
	Ipaddress string
	Port int
	Hostname string
	ConfFile string `toml:"conf-file"`
	IndexFile string `toml:"index-file"`
	DataStoreFile string `toml:"data-file"`
	DebugLevel int `toml:"debug-level"`
	FreezeInterval int `toml:"freeze-interval"`
	FreezeData bool `toml:"freeze-data"`
	LogFile string `toml:"log-file"`
	UseAuth bool `toml:"use-auth"`
	TimeSlew string `toml:"time-slew"`
	TimeSlewDur time.Duration
	ConfRoot string `toml:"conf-root"`
	UseSSL bool `toml:"use-ssl"`
	SslCert string `toml:"ssl-cert"`
	SslKey string `toml:"ssl-key"`
	HttpsUrls bool `toml:"https-urls"`
	DisableWebUI bool `toml:"disable-webui"`
	UseMySQL bool `toml:"use-mysql"`
	MySQL MySQLdb `toml:"mysql"`
	LocalFstoreDir string `toml:"local-filestore-dir"`
}

// MySQL connection options
type MySQLdb struct {
	Username string
	Password string
	Protocol string
	Address string
	Port string
	Dbname string
	ExtraParams map[string]string `toml:"extra_params"`
}

/* Struct for command line options. */
type Options struct {
	Version bool `short:"v" long:"version" description:"Print version info."`
	Verbose []bool `short:"V" long:"verbose" description:"Show verbose debug information. (not implemented)"`
	ConfFile string `short:"c" long:"config" description:"Specify a config file to use."`
	Ipaddress string `short:"I" long:"ipaddress" description:"Listen on a specific IP address."`
	Hostname string `short:"H" long:"hostname" description:"Hostname to use for this server. Defaults to hostname reported by the kernel."`
	Port int `short:"P" long:"port" description:"Port to listen on. If port is set to 443, SSL will be activated. (default: 4545)"`
	IndexFile string `short:"i" long:"index-file" description:"File to save search index data to."`
	DataStoreFile string `short:"D" long:"data-file" description:"File to save data store data to."`
	FreezeInterval int `short:"F" long:"freeze-interval" description:"Interval in seconds to freeze in-memory data structures to disk (requires -i/--index-file and -D/--data-file options to be set). (Default 300 seconds/5 minutes.)"`
	LogFile string `short:"L" long:"log-file" description:"Log to file X"`
	TimeSlew string `long:"time-slew" description:"Time difference allowed between the server's clock at the time in the X-OPS-TIMESTAMP header. Formatted like 5m, 150s, etc. Defaults to 15m."`
	ConfRoot string `long:"conf-root" description:"Root directory for configs and certificates. Default: the directory the config file is in, or the current directory if no config file is set."`
	UseAuth bool `short:"A" long:"use-auth" description:"Use authentication. Default: false."`
	UseSSL bool `long:"use-ssl" description:"Use SSL for connections. If --port is set to 433, this will automatically be turned on. If it is set to 80, it will automatically be turned off. Default: off. Requires --ssl-cert and --ssl-key."`
	SslCert string `long:"ssl-cert" description:"SSL certificate file. If a relative path, will be set relative to --conf-root."`
	SslKey string `long:"ssl-key" description:"SSL key file. If a relative path, will be set relative to --conf-root."`
	HttpsUrls bool `long:"https-urls" description:"Use 'https://' in URLs to server resources if goiardi is not using SSL for its connections. Useful when goiardi is sitting behind a reverse proxy that uses SSL, but is communicating with the proxy over HTTP."`
	DisableWebUI bool `long:"disable-webui" description:"If enabled, disables connections and logins to goiardi over the webui interface."`
	UseMySQL bool `long:"use-mysql" description:"Use a MySQL database for data storage. Configure database options in the config file."`
	LocalFstoreDir string `long:"local-filestore-dir" description:"Directory to save uploaded files in. Optional when running in in-memory mode, *mandatory* for SQL mode."`
}

// The goiardi version.
const Version = "0.4.9999"
// The chef version we're at least aiming for, even if it's not complete yet.
const ChefVersion = "11.0.8"

/* The general plan is to read the command-line options, then parse the config
 * file, fill in the config struct with those values, then apply the
 * command-line options to the config struct. We read the cli options first so
 * we know to look for a different config file if needed, but otherwise the
 * command line options override what's in the config file. */

func InitConfig() *Conf { return &Conf{ } }

// Conf struct with the options specified on the command line or in the config
// file.
var Config = InitConfig()

// Read and apply arguments from the command line.
func ParseConfigOptions() error {
	var opts = &Options{ }
	_, err := flags.Parse(opts)

	if err != nil {
		if err.(*flags.Error).Type == flags.ErrHelp {
			os.Exit(0)
		} else {
			log.Println(err)
			os.Exit(1)
		}
	}

	if opts.Version {
		fmt.Printf("goiardi version %s (aiming for compatibility with Chef Server version %s).\n", Version, ChefVersion)
		os.Exit(0)
	}

	/* Load the config file. Command-line options have precedence over
	 * config file options. */
	if opts.ConfFile != "" {
		if _, err := toml.DecodeFile(opts.ConfFile, Config); err != nil {
			panic(err)
			os.Exit(1)
		}
		Config.ConfFile = opts.ConfFile
		Config.FreezeData = false
	}

	if opts.Hostname != "" {
		Config.Hostname = opts.Hostname
	} else {
		if Config.Hostname == "" {
			Config.Hostname, err = os.Hostname()
			if err != nil {
				log.Println(err)
				Config.Hostname = "localhost"
			}
		}
	}

	if opts.DataStoreFile != "" {
		Config.DataStoreFile = opts.DataStoreFile
	}

	if opts.IndexFile != "" {
		Config.IndexFile = opts.IndexFile
	}

	if !((Config.DataStoreFile == "" && Config.IndexFile == "") || (Config.DataStoreFile != "" && Config.IndexFile != "")) {
		err := fmt.Errorf("-i and -D must either both be specified, or not specified.")
		panic(err)
		os.Exit(1)
	}

	if Config.IndexFile != "" && Config.DataStoreFile != "" {
		Config.FreezeData = true
	}

	if opts.LogFile != "" {
		Config.LogFile = opts.LogFile
	}
	if Config.LogFile != "" {
		lfp, lerr := os.Create(Config.LogFile)
		if lerr != nil {
			log.Println(err)
			os.Exit(1)
		}
		log.SetOutput(lfp)
	}

	/* Database options */
	if opts.UseMySQL {
		Config.UseMySQL = opts.UseMySQL
	}
	// Don't bother setting a default mysql port if mysql isn't used
	if Config.UseMySQL {
		if Config.MySQL.Port == "" {
			Config.MySQL.Port = "3306"
		}
	}

	if opts.LocalFstoreDir != "" {
		Config.LocalFstoreDir = opts.LocalFstoreDir
	}
	if Config.LocalFstoreDir == "" && Config.UseMySQL {
		err := fmt.Errorf("local-filestore-dir must be set when running goiardi in SQL mode")
		log.Println(err)
		os.Exit(1)
	}

	// TODO: once db support is more mature, the dance with freezing data
	// will become more complicated and need to be changed.

	if !Config.FreezeData && (opts.FreezeInterval != 0 || Config.FreezeInterval != 0) {
		log.Printf("FYI, setting the freeze data interval's not especially useful without setting the index and data files.")
	}
	if opts.FreezeInterval != 0 {
		Config.FreezeInterval = opts.FreezeInterval
	}
	if Config.FreezeInterval == 0 {
		Config.FreezeInterval = 300
	}

	/* Root directory for certs and the like */
	if opts.ConfRoot != "" {
		Config.ConfRoot = opts.ConfRoot
	}

	if Config.ConfRoot == "" {
		if Config.ConfFile != "" {
			Config.ConfRoot = path.Dir(Config.ConfFile)
		} else {
			Config.ConfRoot = "."
		}
	}

	Config.Ipaddress = opts.Ipaddress
	if opts.Port != 0 {
		Config.Port = opts.Port
	}
	if Config.Port == 0 {
		Config.Port = 4545
	}

	if opts.UseSSL {
		Config.UseSSL = opts.UseSSL
	}
	if opts.SslCert != "" {
		Config.SslCert = opts.SslCert
	}
	if opts.SslKey != "" {
		Config.SslKey = opts.SslKey
	}
	if opts.HttpsUrls {
		Config.HttpsUrls = opts.HttpsUrls
	}
	// SSL setup
	if Config.Port == 80 {
		Config.UseSSL = false
	} else if Config.Port == 443 {
		Config.UseSSL = true
	}
	if Config.UseSSL {
		if Config.SslCert == "" || Config.SslKey == "" {
			log.Println("SSL mode requires specifying both a certificate and a key file.")
			os.Exit(1)
		}
		/* If the SSL cert and key are not absolute files, join them
		 * with the conf root */
		if !path.IsAbs(Config.SslCert) {
			Config.SslCert = path.Join(Config.ConfRoot, Config.SslCert)
		}
		if !path.IsAbs(Config.SslKey) {
			Config.SslKey = path.Join(Config.ConfRoot, Config.SslKey)
		}
	}

	Config.DebugLevel = len(opts.Verbose)

	if opts.TimeSlew != "" {
		Config.TimeSlew = opts.TimeSlew
	}
	if Config.TimeSlew != "" {
		d, derr := time.ParseDuration(Config.TimeSlew)
		if derr != nil {
			log.Println("Error parsing time-slew:", derr)
			os.Exit(1)
		}
		Config.TimeSlewDur = d
	} else {
		Config.TimeSlewDur, _ = time.ParseDuration("15m")
	}

	if opts.UseAuth {
		Config.UseAuth = opts.UseAuth
	}

	if opts.DisableWebUI {
		Config.DisableWebUI = opts.DisableWebUI
	}

	return nil
}

// The address and port goiardi is configured to listen on.
func ListenAddr() string {
	listen_addr := fmt.Sprintf("%s:%d", Config.Ipaddress, Config.Port)
	return listen_addr
}

// The hostname and port goiardi is configured to use.
func ServerHostname() string {
	var portStr string
	if !(Config.Port == 80 || Config.Port == 443) {
		portStr = fmt.Sprintf(":%d", Config.Port)
	}
	hostname := fmt.Sprintf("%s%s", Config.Hostname, portStr)
	return hostname
}

// The base URL
func ServerBaseURL() string {
	/* TODO: allow configuring using http vs. https */
	var urlScheme string
	if Config.UseSSL || Config.HttpsUrls {
		urlScheme = "https"
	} else {
		urlScheme = "http"
	}
	url := fmt.Sprintf("%s://%s", urlScheme, ServerHostname())
	return url
}
<|endoftext|>


<commit_before>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package config

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/google/syzkaller/fileutil"
	"github.com/google/syzkaller/sys"
	"github.com/google/syzkaller/vm"
)

type Config struct {
	Name     string // Instance name (used for identification and as GCE instance prefix)
	Http     string // TCP address to serve HTTP stats page (e.g. "localhost:50000")
	Rpc      string // TCP address to serve RPC for fuzzer processes (optional, only useful for type "none")
	Workdir  string
	Vmlinux  string
	Kernel   string // e.g. arch/x86/boot/bzImage
	Tag      string // arbitrary optional tag that is saved along with crash reports (e.g. kernel branch/commit)
	Cmdline  string // kernel command line
	Image    string // linux image for VMs
	Initrd   string // linux initial ramdisk. (optional)
	Cpu      int    // number of VM CPUs
	Mem      int    // amount of VM memory in MBs
	Sshkey   string // root ssh key for the image
	Bin      string // qemu/lkvm binary name
	Bin_Args string // additional command line arguments for qemu/lkvm binary
	Debug    bool   // dump all VM output to console
	Output   string // one of stdout/dmesg/file (useful only for local VM)

	Hub_Addr string
	Hub_Key  string

	Syzkaller string   // path to syzkaller checkout (syz-manager will look for binaries in bin subdir)
	Type      string   // VM type (qemu, kvm, local)
	Count     int      // number of VMs (don't secify for adb, instead specify devices)
	Devices   []string // device IDs for adb
	Procs     int      // number of parallel processes inside of every VM

	Sandbox string // type of sandbox to use during fuzzing:
	// "none": don't do anything special (has false positives, e.g. due to killing init)
	// "setuid": impersonate into user nobody (65534), default
	// "namespace": create a new namespace for fuzzer using CLONE_NEWNS/CLONE_NEWNET/CLONE_NEWPID/etc,
	//	requires building kernel with CONFIG_NAMESPACES, CONFIG_UTS_NS, CONFIG_USER_NS, CONFIG_PID_NS and CONFIG_NET_NS.

	Machine_Type string // GCE machine type (e.g. "n1-highcpu-2")

	Cover     bool // use kcov coverage (default: true)
	Leak      bool // do memory leak checking
	Reproduce bool // reproduce, localize and minimize crashers (on by default)

	Enable_Syscalls  []string
	Disable_Syscalls []string
	Suppressions     []string // don't save reports matching these regexps, but reboot VM after them
	Ignores          []string // completely ignore reports matching these regexps (don't save nor reboot)

	// Implementation details beyond this point.
	ParsedSuppressions []*regexp.Regexp `json:"-"`
	ParsedIgnores      []*regexp.Regexp `json:"-"`
}

func Parse(filename string) (*Config, map[int]bool, error) {
	if filename == "" {
		return nil, nil, fmt.Errorf("supply config in -config flag")
	}
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to read config file: %v", err)
	}
	return parse(data)
}

func parse(data []byte) (*Config, map[int]bool, error) {
	unknown, err := checkUnknownFields(data)
	if err != nil {
		return nil, nil, err
	}
	if unknown != "" {
		return nil, nil, fmt.Errorf("unknown field '%v' in config", unknown)
	}
	cfg := new(Config)
	cfg.Cover = true
	cfg.Reproduce = true
	cfg.Sandbox = "setuid"
	if err := json.Unmarshal(data, cfg); err != nil {
		return nil, nil, fmt.Errorf("failed to parse config file: %v", err)
	}
	if _, err := os.Stat(filepath.Join(cfg.Syzkaller, "bin/syz-fuzzer")); err != nil {
		return nil, nil, fmt.Errorf("bad config syzkaller param: can't find bin/syz-fuzzer")
	}
	if _, err := os.Stat(filepath.Join(cfg.Syzkaller, "bin/syz-executor")); err != nil {
		return nil, nil, fmt.Errorf("bad config syzkaller param: can't find bin/syz-executor")
	}
	if cfg.Http == "" {
		return nil, nil, fmt.Errorf("config param http is empty")
	}
	if cfg.Workdir == "" {
		return nil, nil, fmt.Errorf("config param workdir is empty")
	}
	if cfg.Vmlinux == "" {
		return nil, nil, fmt.Errorf("config param vmlinux is empty")
	}
	if cfg.Type == "" {
		return nil, nil, fmt.Errorf("config param type is empty")
	}
	switch cfg.Type {
	case "none":
		if cfg.Count != 0 {
			return nil, nil, fmt.Errorf("invalid config param count: %v, type \"none\" does not support param count", cfg.Count)
		}
		if cfg.Rpc == "" {
			return nil, nil, fmt.Errorf("config param rpc is empty (required for type \"none\")")
		}
		if len(cfg.Devices) != 0 {
			return nil, nil, fmt.Errorf("type %v does not support devices param", cfg.Type)
		}
	case "adb":
		if cfg.Count != 0 {
			return nil, nil, fmt.Errorf("don't specify count for adb, instead specify devices")
		}
		if len(cfg.Devices) == 0 {
			return nil, nil, fmt.Errorf("specify at least 1 adb device")
		}
		cfg.Count = len(cfg.Devices)
	case "gce":
		if cfg.Machine_Type == "" {
			return nil, nil, fmt.Errorf("machine_type parameter is empty (required for gce)")
		}
		fallthrough
	default:
		if cfg.Count <= 0 || cfg.Count > 1000 {
			return nil, nil, fmt.Errorf("invalid config param count: %v, want (1, 1000]", cfg.Count)
		}
		if len(cfg.Devices) != 0 {
			return nil, nil, fmt.Errorf("type %v does not support devices param", cfg.Type)
		}
	}
	if cfg.Rpc == "" {
		cfg.Rpc = "localhost:0"
	}
	if cfg.Procs <= 0 {
		cfg.Procs = 1
	}
	if cfg.Procs > 32 {
		return nil, nil, fmt.Errorf("config param procs has higher value '%v' then the max supported 32", cfg.Procs)
	}
	if cfg.Output == "" {
		if cfg.Type == "local" {
			cfg.Output = "none"
		} else {
			cfg.Output = "stdout"
		}
	}
	switch cfg.Output {
	case "none", "stdout", "dmesg", "file":
	default:
		return nil, nil, fmt.Errorf("config param output must contain one of none/stdout/dmesg/file")
	}
	switch cfg.Sandbox {
	case "none", "setuid", "namespace":
	default:
		return nil, nil, fmt.Errorf("config param sandbox must contain one of none/setuid/namespace")
	}

	wd, err := os.Getwd()
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get wd: %v", err)
	}
	abs := func(path string) string {
		if path != "" && !filepath.IsAbs(path) {
			path = filepath.Join(wd, path)
		}
		return path
	}
	cfg.Workdir = abs(cfg.Workdir)
	cfg.Kernel = abs(cfg.Kernel)
	cfg.Vmlinux = abs(cfg.Vmlinux)
	cfg.Syzkaller = abs(cfg.Syzkaller)
	cfg.Initrd = abs(cfg.Initrd)
	cfg.Sshkey = abs(cfg.Sshkey)
	cfg.Bin = abs(cfg.Bin)

	syscalls, err := parseSyscalls(cfg)
	if err != nil {
		return nil, nil, err
	}

	if err := parseSuppressions(cfg); err != nil {
		return nil, nil, err
	}

	return cfg, syscalls, nil
}

func parseSyscalls(cfg *Config) (map[int]bool, error) {
	match := func(call *sys.Call, str string) bool {
		if str == call.CallName || str == call.Name {
			return true
		}
		if len(str) > 1 && str[len(str)-1] == '*' && strings.HasPrefix(call.Name, str[:len(str)-1]) {
			return true
		}
		return false
	}

	syscalls := make(map[int]bool)
	if len(cfg.Enable_Syscalls) != 0 {
		for _, c := range cfg.Enable_Syscalls {
			n := 0
			for _, call := range sys.Calls {
				if match(call, c) {
					syscalls[call.ID] = true
					n++
				}
			}
			if n == 0 {
				return nil, fmt.Errorf("unknown enabled syscall: %v", c)
			}
		}
	} else {
		for _, call := range sys.Calls {
			syscalls[call.ID] = true
		}
	}
	for _, c := range cfg.Disable_Syscalls {
		n := 0
		for _, call := range sys.Calls {
			if match(call, c) {
				delete(syscalls, call.ID)
				n++
			}
		}
		if n == 0 {
			return nil, fmt.Errorf("unknown disabled syscall: %v", c)
		}
	}
	// mmap is used to allocate memory.
	syscalls[sys.CallMap["mmap"].ID] = true

	return syscalls, nil
}

func parseSuppressions(cfg *Config) error {
	// Add some builtin suppressions.
	supp := append(cfg.Suppressions, []string{
		"panic: failed to start executor binary",
		"panic: executor failed: pthread_create failed",
		"panic: failed to create temp dir",
		"fatal error: runtime: out of memory",
		"fatal error: runtime: cannot allocate memory",
		"fatal error: unexpected signal during runtime execution", // presubmably OOM turned into SIGBUS
		"Out of memory: Kill process .* \\(syz-fuzzer\\)",
		"lowmemorykiller: Killing 'syz-fuzzer'",
		//"INFO: lockdep is turned off", // printed by some sysrq that dumps scheduler state, but also on all lockdep reports
	}...)
	for _, s := range supp {
		re, err := regexp.Compile(s)
		if err != nil {
			return fmt.Errorf("failed to compile suppression '%v': %v", s, err)
		}
		cfg.ParsedSuppressions = append(cfg.ParsedSuppressions, re)
	}
	for _, ignore := range cfg.Ignores {
		re, err := regexp.Compile(ignore)
		if err != nil {
			return fmt.Errorf("failed to compile ignore '%v': %v", ignore, err)
		}
		cfg.ParsedIgnores = append(cfg.ParsedIgnores, re)
	}
	return nil
}

func CreateVMConfig(cfg *Config, index int) (*vm.Config, error) {
	if index < 0 || index >= cfg.Count {
		return nil, fmt.Errorf("invalid VM index %v (count %v)", index, cfg.Count)
	}
	workdir, err := fileutil.ProcessTempDir(cfg.Workdir)
	if err != nil {
		return nil, fmt.Errorf("failed to create instance temp dir: %v", err)
	}
	vmCfg := &vm.Config{
		Name:        fmt.Sprintf("%v-%v-%v", cfg.Type, cfg.Name, index),
		Index:       index,
		Workdir:     workdir,
		Bin:         cfg.Bin,
		BinArgs:     cfg.Bin_Args,
		Kernel:      cfg.Kernel,
		Cmdline:     cfg.Cmdline,
		Image:       cfg.Image,
		Initrd:      cfg.Initrd,
		Sshkey:      cfg.Sshkey,
		Executor:    filepath.Join(cfg.Syzkaller, "bin", "syz-executor"),
		Cpu:         cfg.Cpu,
		Mem:         cfg.Mem,
		Debug:       cfg.Debug,
		MachineType: cfg.Machine_Type,
	}
	if len(cfg.Devices) != 0 {
		vmCfg.Device = cfg.Devices[index]
	}
	return vmCfg, nil
}

func checkUnknownFields(data []byte) (string, error) {
	// While https://github.com/golang/go/issues/15314 is not resolved
	// we don't have a better way than to enumerate all known fields.
	var fields = []string{
		"Name",
		"Http",
		"Rpc",
		"Workdir",
		"Vmlinux",
		"Kernel",
		"Tag",
		"Cmdline",
		"Image",
		"Cpu",
		"Mem",
		"Sshkey",
		"Bin",
		"Bin_Args",
		"Debug",
		"Output",
		"Hub_Addr",
		"Hub_Key",
		"Syzkaller",
		"Type",
		"Count",
		"Devices",
		"Procs",
		"Cover",
		"Reproduce",
		"Sandbox",
		"Leak",
		"Enable_Syscalls",
		"Disable_Syscalls",
		"Suppressions",
		"Ignores",
		"Initrd",
		"Machine_Type",
	}
	f := make(map[string]interface{})
	if err := json.Unmarshal(data, &f); err != nil {
		return "", fmt.Errorf("failed to parse config file: %v", err)
	}
	for k := range f {
		ok := false
		for _, k1 := range fields {
			if strings.ToLower(k) == strings.ToLower(k1) {
				ok = true
				break
			}
		}
		if !ok {
			return k, nil
		}
	}
	return "", nil
}
<commit_msg>config: add another OOM suppression<commit_after>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package config

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/google/syzkaller/fileutil"
	"github.com/google/syzkaller/sys"
	"github.com/google/syzkaller/vm"
)

type Config struct {
	Name     string // Instance name (used for identification and as GCE instance prefix)
	Http     string // TCP address to serve HTTP stats page (e.g. "localhost:50000")
	Rpc      string // TCP address to serve RPC for fuzzer processes (optional, only useful for type "none")
	Workdir  string
	Vmlinux  string
	Kernel   string // e.g. arch/x86/boot/bzImage
	Tag      string // arbitrary optional tag that is saved along with crash reports (e.g. kernel branch/commit)
	Cmdline  string // kernel command line
	Image    string // linux image for VMs
	Initrd   string // linux initial ramdisk. (optional)
	Cpu      int    // number of VM CPUs
	Mem      int    // amount of VM memory in MBs
	Sshkey   string // root ssh key for the image
	Bin      string // qemu/lkvm binary name
	Bin_Args string // additional command line arguments for qemu/lkvm binary
	Debug    bool   // dump all VM output to console
	Output   string // one of stdout/dmesg/file (useful only for local VM)

	Hub_Addr string
	Hub_Key  string

	Syzkaller string   // path to syzkaller checkout (syz-manager will look for binaries in bin subdir)
	Type      string   // VM type (qemu, kvm, local)
	Count     int      // number of VMs (don't secify for adb, instead specify devices)
	Devices   []string // device IDs for adb
	Procs     int      // number of parallel processes inside of every VM

	Sandbox string // type of sandbox to use during fuzzing:
	// "none": don't do anything special (has false positives, e.g. due to killing init)
	// "setuid": impersonate into user nobody (65534), default
	// "namespace": create a new namespace for fuzzer using CLONE_NEWNS/CLONE_NEWNET/CLONE_NEWPID/etc,
	//	requires building kernel with CONFIG_NAMESPACES, CONFIG_UTS_NS, CONFIG_USER_NS, CONFIG_PID_NS and CONFIG_NET_NS.

	Machine_Type string // GCE machine type (e.g. "n1-highcpu-2")

	Cover     bool // use kcov coverage (default: true)
	Leak      bool // do memory leak checking
	Reproduce bool // reproduce, localize and minimize crashers (on by default)

	Enable_Syscalls  []string
	Disable_Syscalls []string
	Suppressions     []string // don't save reports matching these regexps, but reboot VM after them
	Ignores          []string // completely ignore reports matching these regexps (don't save nor reboot)

	// Implementation details beyond this point.
	ParsedSuppressions []*regexp.Regexp `json:"-"`
	ParsedIgnores      []*regexp.Regexp `json:"-"`
}

func Parse(filename string) (*Config, map[int]bool, error) {
	if filename == "" {
		return nil, nil, fmt.Errorf("supply config in -config flag")
	}
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to read config file: %v", err)
	}
	return parse(data)
}

func parse(data []byte) (*Config, map[int]bool, error) {
	unknown, err := checkUnknownFields(data)
	if err != nil {
		return nil, nil, err
	}
	if unknown != "" {
		return nil, nil, fmt.Errorf("unknown field '%v' in config", unknown)
	}
	cfg := new(Config)
	cfg.Cover = true
	cfg.Reproduce = true
	cfg.Sandbox = "setuid"
	if err := json.Unmarshal(data, cfg); err != nil {
		return nil, nil, fmt.Errorf("failed to parse config file: %v", err)
	}
	if _, err := os.Stat(filepath.Join(cfg.Syzkaller, "bin/syz-fuzzer")); err != nil {
		return nil, nil, fmt.Errorf("bad config syzkaller param: can't find bin/syz-fuzzer")
	}
	if _, err := os.Stat(filepath.Join(cfg.Syzkaller, "bin/syz-executor")); err != nil {
		return nil, nil, fmt.Errorf("bad config syzkaller param: can't find bin/syz-executor")
	}
	if cfg.Http == "" {
		return nil, nil, fmt.Errorf("config param http is empty")
	}
	if cfg.Workdir == "" {
		return nil, nil, fmt.Errorf("config param workdir is empty")
	}
	if cfg.Vmlinux == "" {
		return nil, nil, fmt.Errorf("config param vmlinux is empty")
	}
	if cfg.Type == "" {
		return nil, nil, fmt.Errorf("config param type is empty")
	}
	switch cfg.Type {
	case "none":
		if cfg.Count != 0 {
			return nil, nil, fmt.Errorf("invalid config param count: %v, type \"none\" does not support param count", cfg.Count)
		}
		if cfg.Rpc == "" {
			return nil, nil, fmt.Errorf("config param rpc is empty (required for type \"none\")")
		}
		if len(cfg.Devices) != 0 {
			return nil, nil, fmt.Errorf("type %v does not support devices param", cfg.Type)
		}
	case "adb":
		if cfg.Count != 0 {
			return nil, nil, fmt.Errorf("don't specify count for adb, instead specify devices")
		}
		if len(cfg.Devices) == 0 {
			return nil, nil, fmt.Errorf("specify at least 1 adb device")
		}
		cfg.Count = len(cfg.Devices)
	case "gce":
		if cfg.Machine_Type == "" {
			return nil, nil, fmt.Errorf("machine_type parameter is empty (required for gce)")
		}
		fallthrough
	default:
		if cfg.Count <= 0 || cfg.Count > 1000 {
			return nil, nil, fmt.Errorf("invalid config param count: %v, want (1, 1000]", cfg.Count)
		}
		if len(cfg.Devices) != 0 {
			return nil, nil, fmt.Errorf("type %v does not support devices param", cfg.Type)
		}
	}
	if cfg.Rpc == "" {
		cfg.Rpc = "localhost:0"
	}
	if cfg.Procs <= 0 {
		cfg.Procs = 1
	}
	if cfg.Procs > 32 {
		return nil, nil, fmt.Errorf("config param procs has higher value '%v' then the max supported 32", cfg.Procs)
	}
	if cfg.Output == "" {
		if cfg.Type == "local" {
			cfg.Output = "none"
		} else {
			cfg.Output = "stdout"
		}
	}
	switch cfg.Output {
	case "none", "stdout", "dmesg", "file":
	default:
		return nil, nil, fmt.Errorf("config param output must contain one of none/stdout/dmesg/file")
	}
	switch cfg.Sandbox {
	case "none", "setuid", "namespace":
	default:
		return nil, nil, fmt.Errorf("config param sandbox must contain one of none/setuid/namespace")
	}

	wd, err := os.Getwd()
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get wd: %v", err)
	}
	abs := func(path string) string {
		if path != "" && !filepath.IsAbs(path) {
			path = filepath.Join(wd, path)
		}
		return path
	}
	cfg.Workdir = abs(cfg.Workdir)
	cfg.Kernel = abs(cfg.Kernel)
	cfg.Vmlinux = abs(cfg.Vmlinux)
	cfg.Syzkaller = abs(cfg.Syzkaller)
	cfg.Initrd = abs(cfg.Initrd)
	cfg.Sshkey = abs(cfg.Sshkey)
	cfg.Bin = abs(cfg.Bin)

	syscalls, err := parseSyscalls(cfg)
	if err != nil {
		return nil, nil, err
	}

	if err := parseSuppressions(cfg); err != nil {
		return nil, nil, err
	}

	return cfg, syscalls, nil
}

func parseSyscalls(cfg *Config) (map[int]bool, error) {
	match := func(call *sys.Call, str string) bool {
		if str == call.CallName || str == call.Name {
			return true
		}
		if len(str) > 1 && str[len(str)-1] == '*' && strings.HasPrefix(call.Name, str[:len(str)-1]) {
			return true
		}
		return false
	}

	syscalls := make(map[int]bool)
	if len(cfg.Enable_Syscalls) != 0 {
		for _, c := range cfg.Enable_Syscalls {
			n := 0
			for _, call := range sys.Calls {
				if match(call, c) {
					syscalls[call.ID] = true
					n++
				}
			}
			if n == 0 {
				return nil, fmt.Errorf("unknown enabled syscall: %v", c)
			}
		}
	} else {
		for _, call := range sys.Calls {
			syscalls[call.ID] = true
		}
	}
	for _, c := range cfg.Disable_Syscalls {
		n := 0
		for _, call := range sys.Calls {
			if match(call, c) {
				delete(syscalls, call.ID)
				n++
			}
		}
		if n == 0 {
			return nil, fmt.Errorf("unknown disabled syscall: %v", c)
		}
	}
	// mmap is used to allocate memory.
	syscalls[sys.CallMap["mmap"].ID] = true

	return syscalls, nil
}

func parseSuppressions(cfg *Config) error {
	// Add some builtin suppressions.
	supp := append(cfg.Suppressions, []string{
		"panic: failed to start executor binary",
		"panic: executor failed: pthread_create failed",
		"panic: failed to create temp dir",
		"fatal error: runtime: out of memory",
		"fatal error: runtime: cannot allocate memory",
		"fatal error: unexpected signal during runtime execution", // presubmably OOM turned into SIGBUS
		"signal SIGBUS: bus error",                                // presubmably OOM turned into SIGBUS
		"Out of memory: Kill process .* \\(syz-fuzzer\\)",
		"lowmemorykiller: Killing 'syz-fuzzer'",
		//"INFO: lockdep is turned off", // printed by some sysrq that dumps scheduler state, but also on all lockdep reports
	}...)
	for _, s := range supp {
		re, err := regexp.Compile(s)
		if err != nil {
			return fmt.Errorf("failed to compile suppression '%v': %v", s, err)
		}
		cfg.ParsedSuppressions = append(cfg.ParsedSuppressions, re)
	}
	for _, ignore := range cfg.Ignores {
		re, err := regexp.Compile(ignore)
		if err != nil {
			return fmt.Errorf("failed to compile ignore '%v': %v", ignore, err)
		}
		cfg.ParsedIgnores = append(cfg.ParsedIgnores, re)
	}
	return nil
}

func CreateVMConfig(cfg *Config, index int) (*vm.Config, error) {
	if index < 0 || index >= cfg.Count {
		return nil, fmt.Errorf("invalid VM index %v (count %v)", index, cfg.Count)
	}
	workdir, err := fileutil.ProcessTempDir(cfg.Workdir)
	if err != nil {
		return nil, fmt.Errorf("failed to create instance temp dir: %v", err)
	}
	vmCfg := &vm.Config{
		Name:        fmt.Sprintf("%v-%v-%v", cfg.Type, cfg.Name, index),
		Index:       index,
		Workdir:     workdir,
		Bin:         cfg.Bin,
		BinArgs:     cfg.Bin_Args,
		Kernel:      cfg.Kernel,
		Cmdline:     cfg.Cmdline,
		Image:       cfg.Image,
		Initrd:      cfg.Initrd,
		Sshkey:      cfg.Sshkey,
		Executor:    filepath.Join(cfg.Syzkaller, "bin", "syz-executor"),
		Cpu:         cfg.Cpu,
		Mem:         cfg.Mem,
		Debug:       cfg.Debug,
		MachineType: cfg.Machine_Type,
	}
	if len(cfg.Devices) != 0 {
		vmCfg.Device = cfg.Devices[index]
	}
	return vmCfg, nil
}

func checkUnknownFields(data []byte) (string, error) {
	// While https://github.com/golang/go/issues/15314 is not resolved
	// we don't have a better way than to enumerate all known fields.
	var fields = []string{
		"Name",
		"Http",
		"Rpc",
		"Workdir",
		"Vmlinux",
		"Kernel",
		"Tag",
		"Cmdline",
		"Image",
		"Cpu",
		"Mem",
		"Sshkey",
		"Bin",
		"Bin_Args",
		"Debug",
		"Output",
		"Hub_Addr",
		"Hub_Key",
		"Syzkaller",
		"Type",
		"Count",
		"Devices",
		"Procs",
		"Cover",
		"Reproduce",
		"Sandbox",
		"Leak",
		"Enable_Syscalls",
		"Disable_Syscalls",
		"Suppressions",
		"Ignores",
		"Initrd",
		"Machine_Type",
	}
	f := make(map[string]interface{})
	if err := json.Unmarshal(data, &f); err != nil {
		return "", fmt.Errorf("failed to parse config file: %v", err)
	}
	for k := range f {
		ok := false
		for _, k1 := range fields {
			if strings.ToLower(k) == strings.ToLower(k1) {
				ok = true
				break
			}
		}
		if !ok {
			return k, nil
		}
	}
	return "", nil
}
<|endoftext|>


<commit_before>// io.Reader --> map[string]interface{} or JSON string
// nothing magic - just implements generic Go case

package x2j

import (
	"encoding/json"
	"encoding/xml"
	"io"
)

// ToTree() - parse a XML io.Reader to a tree of Nodes
func ToTree(rdr io.Reader) (*Node, error) {
	// We need to put an *os.File reader in a ByteReader or the xml.NewDecoder
	// will wrap it in a bufio.Reader and seek on the file beyond where the
	// xml.Decoder parses!
	if _, ok := rdr.(io.ByteReader); !ok {
		rdr = myByteReader(rdr) // see code at EOF
	}

	p := xml.NewDecoder(rdr)
	p.CharsetReader = X2jCharsetReader
	n, perr := xmlToTree("", nil, p)
	if perr != nil {
		return nil, perr
	}

	return n, nil
}

// ToMap() - parse a XML io.Reader to a map[string]interface{}
func ToMap(rdr io.Reader, recast ...bool) (map[string]interface{}, error) {
	var r bool
	if len(recast) == 1 {
		r = recast[0]
	}
	n, err := ToTree(rdr)
	if err != nil {
		return nil, err
	}

	m := make(map[string]interface{})
	m[n.key] = n.treeToMap(r)

	return m, nil
}

// ToJson() - parse a XML io.Reader to a JSON string
func ToJson(rdr io.Reader, recast ...bool) (string, error) {
	var r bool
	if len(recast) == 1 {
		r = recast[0]
	}
	m, merr := ToMap(rdr, r)
	if m == nil || merr != nil {
		return "", merr
	}

	b, berr := json.Marshal(m)
	if berr != nil {
		return "", berr
	}

	return string(b), nil
}

// ToJsonIndent - the pretty form of ReaderToJson
func ToJsonIndent(rdr io.Reader, recast ...bool) (string, error) {
	var r bool
	if len(recast) == 1 {
		r = recast[0]
	}
	m, merr := ToMap(rdr, r)
	if m == nil || merr != nil {
		return "", merr
	}

	b, berr := json.MarshalIndent(m, "", "  ")
	if berr != nil {
		return "", berr
	}

	// NOTE: don't have to worry about safe JSON marshaling with json.Marshal, since '<' and '>" are reservedin XML.
	return string(b), nil
}

// ReaderValuesFromTagPath - io.Reader version of ValuesFromTagPath()
func ReaderValuesFromTagPath(rdr io.Reader, path string, getAttrs ...bool) ([]interface{}, error) {
	var a bool
	if len(getAttrs) == 1 {
		a = getAttrs[0]
	}
	m, err := ToMap(rdr)
	if err != nil {
		return nil, err
	}

	return ValuesFromKeyPath(m, path, a), nil
}

// ReaderValuesForTag - io.Reader version of ValuesForTag()
func ReaderValuesForTag(rdr io.Reader, tag string) ([]interface{}, error) {
	m, err := ToMap(rdr)
	if err != nil {
		return nil, err
	}

	return ValuesForKey(m, tag), nil
}

//============================ from github.com/clbanning/mxj/mxl.go ==========================

type byteReader struct {
	r io.Reader
	b []byte
}

func myByteReader(r io.Reader) io.Reader {
	b := make([]byte, 1)
	return &byteReader{r, b}
}

// need for io.Reader - but we don't use it ...
func (b *byteReader) Read(p []byte) (int, error) {
	return 0, nil
}

func (b *byteReader) ReadByte() (byte, error) {
	n, err := b.r.Read(b.b)
	if n == 1 {
		return b.b[0], err
	} else {
		var c byte
		return c, err
	}
}

<commit_msg>handle net/http.Response.Body *io.Reader error returns<commit_after>// io.Reader --> map[string]interface{} or JSON string
// nothing magic - just implements generic Go case

package x2j

import (
	"encoding/json"
	"encoding/xml"
	"io"
)

// ToTree() - parse a XML io.Reader to a tree of Nodes
func ToTree(rdr io.Reader) (*Node, error) {
	// We need to put an *os.File reader in a ByteReader or the xml.NewDecoder
	// will wrap it in a bufio.Reader and seek on the file beyond where the
	// xml.Decoder parses!
	if _, ok := rdr.(io.ByteReader); !ok {
		rdr = myByteReader(rdr) // see code at EOF
	}

	p := xml.NewDecoder(rdr)
	p.CharsetReader = X2jCharsetReader
	n, perr := xmlToTree("", nil, p)
	if perr != nil {
		return nil, perr
	}

	return n, nil
}

// ToMap() - parse a XML io.Reader to a map[string]interface{}
func ToMap(rdr io.Reader, recast ...bool) (map[string]interface{}, error) {
	var r bool
	if len(recast) == 1 {
		r = recast[0]
	}
	n, err := ToTree(rdr)
	if err != nil {
		return nil, err
	}

	m := make(map[string]interface{})
	m[n.key] = n.treeToMap(r)

	return m, nil
}

// ToJson() - parse a XML io.Reader to a JSON string
func ToJson(rdr io.Reader, recast ...bool) (string, error) {
	var r bool
	if len(recast) == 1 {
		r = recast[0]
	}
	m, merr := ToMap(rdr, r)
	if m == nil || merr != nil {
		return "", merr
	}

	b, berr := json.Marshal(m)
	if berr != nil {
		return "", berr
	}

	return string(b), nil
}

// ToJsonIndent - the pretty form of ReaderToJson
func ToJsonIndent(rdr io.Reader, recast ...bool) (string, error) {
	var r bool
	if len(recast) == 1 {
		r = recast[0]
	}
	m, merr := ToMap(rdr, r)
	if m == nil || merr != nil {
		return "", merr
	}

	b, berr := json.MarshalIndent(m, "", "  ")
	if berr != nil {
		return "", berr
	}

	// NOTE: don't have to worry about safe JSON marshaling with json.Marshal, since '<' and '>" are reservedin XML.
	return string(b), nil
}

// ReaderValuesFromTagPath - io.Reader version of ValuesFromTagPath()
func ReaderValuesFromTagPath(rdr io.Reader, path string, getAttrs ...bool) ([]interface{}, error) {
	var a bool
	if len(getAttrs) == 1 {
		a = getAttrs[0]
	}
	m, err := ToMap(rdr)
	if err != nil {
		return nil, err
	}

	return ValuesFromKeyPath(m, path, a), nil
}

// ReaderValuesForTag - io.Reader version of ValuesForTag()
func ReaderValuesForTag(rdr io.Reader, tag string) ([]interface{}, error) {
	m, err := ToMap(rdr)
	if err != nil {
		return nil, err
	}

	return ValuesForKey(m, tag), nil
}

//============================ from github.com/clbanning/mxj/mxl.go ==========================

type byteReader struct {
	r io.Reader
	b []byte
}

func myByteReader(r io.Reader) io.Reader {
	b := make([]byte, 1)
	return &byteReader{r, b}
}

// need for io.Reader - but we don't use it ...
func (b *byteReader) Read(p []byte) (int, error) {
	return 0, nil
}

func (b *byteReader) ReadByte() (byte, error) {
	_, err := b.r.Read(b.b)
	if len(b.b) > 0 {
		return b.b[0], nil
	} else {
		var c byte
		return c, err
	}
}

<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"net"
	"strconv"

	log "github.com/cihub/seelog"
	"github.com/stampzilla/stampzilla-go/nodes/stampzilla-server/logic"
	"github.com/stampzilla/stampzilla-go/nodes/stampzilla-server/metrics"
	serverprotocol "github.com/stampzilla/stampzilla-go/nodes/stampzilla-server/protocol"
)

type NodeServer struct {
	Config           *ServerConfig         `inject:""`
	Logic            *logic.Logic          `inject:""`
	Nodes            *serverprotocol.Nodes `inject:""`
	WebsocketHandler *WebsocketHandler     `inject:""`
	ElasticSearch    *ElasticSearch        `inject:""`
	Metrics          *metrics.Metrics      `inject:""`

	ServerNode   *serverprotocol.Node
	State        map[string]interface{}
	logicChannel chan string
}

func NewNodeServer() *NodeServer {
	return &NodeServer{
		ServerNode: &serverprotocol.Node{},
		State:      make(map[string]interface{}),
	}
}

func (ns *NodeServer) Start() {
	log.Info("Starting NodeServer (:" + ns.Config.NodePort + ")")
	listen, err := net.Listen("tcp", ":"+ns.Config.NodePort)
	if err != nil {
		log.Error("listen error", err)
		return
	}

	ns.Logic.RestoreRulesFromFile("rules.json")
	ns.addServerNode()

	ns.logicChannel = ns.Logic.ListenForChanges(ns.ServerNode.Uuid)

	//ns.Logic.Listen()

	//return
	go func() {
		for {
			fd, err := listen.Accept()
			if err != nil {
				log.Error("accept error", err)
				return
			}

			go ns.newNodeConnection(fd)
		}
	}()
}

func (ns *NodeServer) newNodeConnection(connection net.Conn) {
	// Recive data
	log.Info("New client connected")
	name := ""
	uuid := ""
	decoder := json.NewDecoder(connection)
	//encoder := json.NewEncoder(os.Stdout)
	var logicChannel chan string
	for {
		var node serverprotocol.Node
		err := decoder.Decode(&node)

		if err != nil {
			if err.Error() == "EOF" {
				log.Info(name, " - Client disconnected")
				if uuid != "" {
					ns.Nodes.Delete(uuid)
					close(logicChannel)
				}
				//TODO be able to not send everything always. perhaps implement remove instead of all?
				ns.WebsocketHandler.SendAllNodes()
				return
			}
			log.Warn("Not disconnect but error: ", err)
			//return here?
		} else {
			name = node.Name
			uuid = node.Uuid

			existingNode := ns.Nodes.ByUuid(uuid)
			if existingNode == nil {
				ns.Nodes.Add(&node)
				logicChannel = ns.Logic.ListenForChanges(node.Uuid)
				node.SetConn(connection)
				ns.updateState(logicChannel, &node)
			} else {
				existingNode.SetState(node.State())
				ns.updateState(logicChannel, existingNode)
			}

			ns.WebsocketHandler.SendSingleNode(uuid)
		}
	}
}

func (ns *NodeServer) updateState(updateChan chan string, node *serverprotocol.Node) {
	if node == nil {
		log.Warn("Recived an updateState but no node was provided, ignoring...")
		return
	}

	//log.Info(node.Uuid, " - ", node.Name, " - Got update on state")

	//Send to logic for evaluation
	//logicChannel := ns.Logic.ListenForChanges(node.Uuid)
	//state, _ := json.Marshal(node.State)
	//*logicChannel <- string(state)
	ns.Logic.Update(updateChan, node)

	//Send to metrics
	ns.Metrics.Update(node)
}

func (self *NodeServer) addServerNode() {
	self.ServerNode.Name = "server"
	self.ServerNode.Uuid = self.Config.Uuid
	self.ServerNode.SetState(self.State)

	self.Nodes.Add(self.ServerNode)
}

func (self *NodeServer) Set(key string, value interface{}) {
	self.State[key] = cast(value)
	self.updateState(self.logicChannel, self.ServerNode)
}

func (self *NodeServer) Trigger(key string, value interface{}) {
	self.Set(key, value)

	switch self.State[key].(type) {
	case int:
		self.Set(key, 0)
	case float64:
		self.Set(key, 0.0)
	case string:
		self.Set(key, "")
	case bool:
		self.Set(key, 0)
	}
}

func cast(s interface{}) interface{} {
	switch v := s.(type) {
	case int:
		return v
		//return strconv.Itoa(v)
	case float64:
		//return strconv.FormatFloat(v, 'f', -1, 64)
		return v
	case string:
		if n, err := strconv.Atoi(v); err == nil {
			return n
		}
		return v
	case bool:
		if v {
			return 1
		}
		return 0
	}
	return ""
}
<commit_msg>Tried to fix the disconnection bug, lets se if it helps<commit_after>package main

import (
	"encoding/json"
	"net"
	"strconv"

	log "github.com/cihub/seelog"
	"github.com/stampzilla/stampzilla-go/nodes/stampzilla-server/logic"
	"github.com/stampzilla/stampzilla-go/nodes/stampzilla-server/metrics"
	serverprotocol "github.com/stampzilla/stampzilla-go/nodes/stampzilla-server/protocol"
)

type NodeServer struct {
	Config           *ServerConfig         `inject:""`
	Logic            *logic.Logic          `inject:""`
	Nodes            *serverprotocol.Nodes `inject:""`
	WebsocketHandler *WebsocketHandler     `inject:""`
	ElasticSearch    *ElasticSearch        `inject:""`
	Metrics          *metrics.Metrics      `inject:""`

	ServerNode   *serverprotocol.Node
	State        map[string]interface{}
	logicChannel chan string
}

func NewNodeServer() *NodeServer {
	return &NodeServer{
		ServerNode: &serverprotocol.Node{},
		State:      make(map[string]interface{}),
	}
}

func (ns *NodeServer) Start() {
	log.Info("Starting NodeServer (:" + ns.Config.NodePort + ")")
	listen, err := net.Listen("tcp", ":"+ns.Config.NodePort)
	if err != nil {
		log.Error("listen error", err)
		return
	}

	ns.Logic.RestoreRulesFromFile("rules.json")
	ns.addServerNode()

	ns.logicChannel = ns.Logic.ListenForChanges(ns.ServerNode.Uuid)

	//ns.Logic.Listen()

	//return
	go func() {
		for {
			fd, err := listen.Accept()
			if err != nil {
				log.Error("accept error", err)
				return
			}

			go ns.newNodeConnection(fd)
		}
	}()
}

func (ns *NodeServer) newNodeConnection(connection net.Conn) {
	// Recive data
	log.Info("New client connected (", connection.RemoteAddr(), ")")
	name := ""
	uuid := ""
	decoder := json.NewDecoder(connection)
	//encoder := json.NewEncoder(os.Stdout)
	var logicChannel chan string
	for {
		var node serverprotocol.Node
		err := decoder.Decode(&node)

		if err != nil {
			if err.Error() == "EOF" {
				log.Info(name, " - Client disconnected")
				if uuid != "" {
					ns.Nodes.Delete(uuid)
					close(logicChannel)
				}
				//TODO be able to not send everything always. perhaps implement remove instead of all?
				ns.WebsocketHandler.SendAllNodes()
				return
			}
			log.Warn(name, " - Not disconnect but error (closing connection): ", err)

			connection.Close()
			return
		} else {
			name = node.Name
			uuid = node.Uuid

			existingNode := ns.Nodes.ByUuid(uuid)
			if existingNode == nil {
				ns.Nodes.Add(&node)
				logicChannel = ns.Logic.ListenForChanges(node.Uuid)
				node.SetConn(connection)
				ns.updateState(logicChannel, &node)
			} else {
				existingNode.SetState(node.State())
				ns.updateState(logicChannel, existingNode)
			}

			ns.WebsocketHandler.SendSingleNode(uuid)
		}
	}
}

func (ns *NodeServer) updateState(updateChan chan string, node *serverprotocol.Node) {
	if node == nil {
		log.Warn("Recived an updateState but no node was provided, ignoring...")
		return
	}

	//log.Info(node.Uuid, " - ", node.Name, " - Got update on state")

	//Send to logic for evaluation
	//logicChannel := ns.Logic.ListenForChanges(node.Uuid)
	//state, _ := json.Marshal(node.State)
	//*logicChannel <- string(state)
	ns.Logic.Update(updateChan, node)

	//Send to metrics
	ns.Metrics.Update(node)
}

func (self *NodeServer) addServerNode() {
	self.ServerNode.Name = "server"
	self.ServerNode.Uuid = self.Config.Uuid
	self.ServerNode.SetState(self.State)

	self.Nodes.Add(self.ServerNode)
}

func (self *NodeServer) Set(key string, value interface{}) {
	self.State[key] = cast(value)
	self.updateState(self.logicChannel, self.ServerNode)
}

func (self *NodeServer) Trigger(key string, value interface{}) {
	self.Set(key, value)

	switch self.State[key].(type) {
	case int:
		self.Set(key, 0)
	case float64:
		self.Set(key, 0.0)
	case string:
		self.Set(key, "")
	case bool:
		self.Set(key, 0)
	}
}

func cast(s interface{}) interface{} {
	switch v := s.(type) {
	case int:
		return v
		//return strconv.Itoa(v)
	case float64:
		//return strconv.FormatFloat(v, 'f', -1, 64)
		return v
	case string:
		if n, err := strconv.Atoi(v); err == nil {
			return n
		}
		return v
	case bool:
		if v {
			return 1
		}
		return 0
	}
	return ""
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/gin-gonic/gin"
)

const VERSION = "0.1.1"

type Response struct {
	Success bool   `json:"success"`
	Message string `json:"message"`
}

type StatusResponse struct {
	Running       bool   `json:"running"`            // True if player is running
	File          string `json:"file"`               // Path to current media file
	Name          string `json:"name"`               // Titleized filename
	MovieDbApiKey string `json:"tmdb_api_key"`       // TheMovieDb API key
	Position      string `json:"position,omitempty"` // Current position in the movie
	Duration      string `json:"duration,omitempty"` // Movie duration
}

type FileEntry struct {
	Filename string `json:"filename"`
	IsDir    bool   `json:"directory"`
}

var (
	// Regular expression to match all supported video files
	RegexFormats = regexp.MustCompile(`.(avi|mpg|mov|flv|wmv|asf|mpeg|m4v|divx|mp4|mkv)$`)

	// Regular expression to convert filenames to titles
	RegexBrackets = regexp.MustCompile(`[\(\[\]\)]`)
	RegexYear     = regexp.MustCompile(`((19|20)[\d]{2})`)
	RegexEpisode  = regexp.MustCompile(`(?i)S[\d]+E[\d]+`)
	RegexJunk     = regexp.MustCompile(`(?i)(1080p|720p|3d|brrip|bluray|webrip|x264|aac)`)
	RegexSpace    = regexp.MustCompile(`\s{2,}`)

	// OMXPlayer control commands, these are piped via STDIN to omxplayer process
	Commands = map[string]string{
		"pause":             "p",            // Pause/continue playback
		"stop":              "q",            // Stop playback and exit
		"volume_up":         "+",            // Change volume by +3dB
		"volume_down":       "-",            // Change volume by -3dB
		"subtitles":         "s",            // Enable/disable subtitles
		"seek_back":         "\x1b\x5b\x44", // Seek -30 seconds
		"seek_back_fast":    "\x1b\x5b\x42", // Seek -600 second
		"seek_forward":      "\x1b\x5b\x43", // Seek +30 second
		"seek_forward_fast": "\x1b\x5b\x41", // Seek +600 seconds
	}

	// Path where all media files are stored
	MediaPath string

	// Path to omxplayer executable
	OmxPath string

	// Child process for spawning omxplayer
	Omx *exec.Cmd

	// Child process STDIN pipe to send commands
	OmxIn io.WriteCloser

	// Channel to pass along commands to the player routine
	Command chan string

	// Currently playing media file name
	CurrentFile string

	// Enable Zeroconf discovery
	Zeroconf bool

	// Current stream
	stream *Stream
)

// Returns true if specified file exists
func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

// Scan given path for all directories and matching video files.
// If nothing was found it will return an empty slice.
func scanPath(path string) []FileEntry {
	entries := make([]FileEntry, 0)

	files, err := ioutil.ReadDir(path)
	if err != nil {
		return entries
	}

	for _, file := range files {
		entry := FileEntry{
			Filename: file.Name(),
			IsDir:    file.IsDir(),
		}

		// Do not include non-video files in the list
		if !file.IsDir() && !omxCanPlay(file.Name()) {
			continue
		}

		entries = append(entries, entry)
	}

	return entries
}

// Convert media filename to regular title
func fileToTitle(name string) string {
	// Remove file extension from name
	name = strings.Replace(name, filepath.Ext(name), "", 1)

	// Replace all dots with white space
	name = strings.Replace(name, ".", " ", -1)

	// Replace parenteses and brackets
	name = RegexBrackets.ReplaceAllString(name, "")

	// Check if name has a typical torrent format: "name year format etc"
	pos := RegexYear.FindStringIndex(name)
	if len(pos) > 0 {
		name = name[0:pos[0]]
	} else {
		// Check if its an episode, i.e. "S##E##"
		pos = RegexEpisode.FindStringIndex(name)

		if len(pos) > 0 {
			name = name[0:pos[0]]
		}
	}

	// Remove junk stuff
	name = RegexJunk.ReplaceAllString(name, "")

	// Remove any extra white space
	name = RegexSpace.ReplaceAllString(name, "")

	// Truncate space
	name = strings.TrimSpace(name)

	return name
}

// Determine the full path to omxplayer executable. Returns error if not found.
func omxDetect() error {
	buff, err := exec.Command("which", "omxplayer").Output()
	if err != nil {
		return err
	}

	// Set path in global variable
	OmxPath = strings.TrimSpace(string(buff))

	return nil
}

// Start command listener. Commands are coming in through a channel.
func omxListen() {
	Command = make(chan string)

	for {
		command := <-Command

		// Skip command handling of omx player is not active
		if Omx == nil {
			continue
		}

		// Send command to the player
		omxWrite(command)

		// Attempt to kill the process if stop command is requested
		if command == "stop" {
			Omx.Process.Kill()
		}
	}
}

// Start omxplayer playback for a given video file. Returns error if start fails.
func omxPlay(file string) error {
	Omx = exec.Command(
		OmxPath,       // path to omxplayer executable
		"--stats",     // print stats to stdout (buffers, time, etc)
		"--with-info", // print stats about streams before playback
		"--refresh",   // adjust framerate/resolution to video
		"--blank",     // set background to black
		"--adev",      // audio out device
		"hdmi",        // using hdmi for audio/video
		file,          // path to video file
	)

	// Grab child process STDIN
	stdin, err := Omx.StdinPipe()
	if err != nil {
		return err
	}
	defer stdin.Close()

	stderr, err := Omx.StderrPipe()
	if err != nil {
		return err
	}
	defer stderr.Close()

	stdout, err := Omx.StdoutPipe()
	if err != nil {
		return err
	}
	defer stdout.Close()

	stream = NewStream()
	go stream.Start(stdout, stderr)

	// Start omxplayer execution.
	// If successful, something will appear on HDMI display.
	err = Omx.Start()
	if err != nil {
		return err
	}

	// Set current file
	CurrentFile = file

	// Make child's STDIN globally available
	OmxIn = stdin

	// Wait until child process is finished
	err = Omx.Wait()
	if err != nil {
		fmt.Fprintln(os.Stdout, "Process exited with error:", err)
	}

	omxCleanup()

	return nil
}

// Write a command string to the omxplayer process's STDIN
func omxWrite(command string) {
	if OmxIn != nil {
		io.WriteString(OmxIn, Commands[command])
	}
}

// Terminate any running omxplayer processes. Fixes random hangs.
func omxKill() {
	exec.Command("killall", "omxplayer.bin").Output()
	exec.Command("killall", "omxplayer").Output()
}

// Reset internal state and stop any running processes
func omxCleanup() {
	Omx = nil
	OmxIn = nil
	CurrentFile = ""

	omxKill()

	stream = nil
}

// Check if player is currently active
func omxIsActive() bool {
	return Omx != nil
}

// Check if player can play the file
func omxCanPlay(path string) bool {
	if RegexFormats.Match([]byte(path)) {
		return true
	}

	return false
}

// Check if TheMovieDB API could be used
func movieDbApiKey() string {
	return os.Getenv("TMDB_API_KEY")
}

func httpBrowse(c *gin.Context) {
	path := c.Request.FormValue("path")

	if path != "" {
		path = fmt.Sprintf("%s/%s", MediaPath, path)
	} else {
		path = MediaPath
	}

	c.JSON(200, scanPath(path))
}

func httpCommand(c *gin.Context) {
	val := c.Params.ByName("command")

	if _, ok := Commands[val]; !ok {
		c.JSON(400, Response{false, "Invalid command"})
		return
	}

	fmt.Println("Received command:", val)

	// Handle requested commmand
	Command <- val

	c.JSON(200, Response{true, "OK"})
}

func httpServe(c *gin.Context) {
	file := c.Request.URL.Query().Get("file")
	if file == "" {
		return
	}

	file = fmt.Sprintf("%s/%s", MediaPath, file)
	if !fileExists(file) {
		c.String(404, "Not found")
		return
	}

	if !omxCanPlay(file) {
		c.String(400, "Invalid format")
		return
	}

	http.ServeFile(c.Writer, c.Request, file)
}

func httpPlay(c *gin.Context) {
	if omxIsActive() {
		c.JSON(400, Response{false, "Player is already running"})
		return
	}

	file := c.Request.FormValue("file")
	if file == "" {
		c.JSON(400, Response{false, "File is required"})
		return
	}

	file = fmt.Sprintf("%s/%s", MediaPath, file)

	if !fileExists(file) {
		c.JSON(400, Response{false, "File does not exist"})
		return
	}

	if !omxCanPlay(file) {
		c.JSON(400, Response{false, "File cannot be played"})
		return
	}

	go omxPlay(file)

	c.JSON(200, Response{true, "OK"})
}

func httpStatus(c *gin.Context) {
	resp := StatusResponse{
		Running:       omxIsActive(),
		File:          CurrentFile,
		Name:          fileToTitle(filepath.Base(CurrentFile)),
		MovieDbApiKey: movieDbApiKey(),
	}

	if stream != nil {
		resp.Duration = durationFromSeconds(stream.duration)
		resp.Position = stream.pos.String()
	}

	c.JSON(200, resp)
}

func httpIndex(c *gin.Context) {
	data, err := Asset("static/index.html")

	if err != nil {
		c.String(400, err.Error())
		return
	}

	c.Data(200, "text/html; charset=utf-8", data)
}

func terminate(message string, code int) {
	fmt.Println(message)
	os.Exit(code)
}

func usage() {
	terminate("Usage: omxremote path/to/media/dir", 0)
}

func init() {
	var printVersion bool

	flag.StringVar(&MediaPath, "media", "./", "Path to media files")
	flag.BoolVar(&Zeroconf, "zeroconf", false, "Enable service advertisement with Zeroconf")
	flag.BoolVar(&printVersion, "v", false, "Print version")
	flag.Parse()

	if printVersion {
		fmt.Printf("omxremote v%v\n", VERSION)
		os.Exit(0)
	}
}

func main() {
	// Expand media path if needed
	MediaPath = strings.Replace(MediaPath, "~", os.Getenv("HOME"), 1)

	// Get path from arguments and remove trailing slash
	MediaPath = strings.TrimRight(MediaPath, "/")

	if !fileExists(MediaPath) {
		terminate(fmt.Sprintf("Directory does not exist: %s", MediaPath), 1)
	}

	// Check if player is installed
	if omxDetect() != nil {
		terminate("omxplayer is not installed", 1)
	}

	// Make sure nothing is running
	omxCleanup()

	// Start a remote command listener
	go omxListen()

	// Start zeroconf service advertisement
	stopZeroconf := make(chan bool)
	go startZeroConfAdvertisement(stopZeroconf)

	// Disable debugging mode
	gin.SetMode("release")

	// Setup HTTP server
	router := gin.Default()

	// Handle CORS
	router.Use(func(c *gin.Context) {
		c.Header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
		c.Header("Access-Control-Allow-Origin", "*")
		c.Header("Access-Control-Expose-Headers", "*")
	})

	router.GET("/", httpIndex)
	router.GET("/status", httpStatus)
	router.GET("/browse", httpBrowse)
	router.GET("/play", httpPlay)
	router.GET("/serve", httpServe)
	router.GET("/command/:command", httpCommand)

	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}

	fmt.Println("Starting server on 0.0.0.0:" + port)
	router.Run(":" + port)
}
<commit_msg>Fix #3: starts zeroconf only if requested<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/gin-gonic/gin"
)

const VERSION = "0.1.1"

type Response struct {
	Success bool   `json:"success"`
	Message string `json:"message"`
}

type StatusResponse struct {
	Running       bool   `json:"running"`            // True if player is running
	File          string `json:"file"`               // Path to current media file
	Name          string `json:"name"`               // Titleized filename
	MovieDbApiKey string `json:"tmdb_api_key"`       // TheMovieDb API key
	Position      string `json:"position,omitempty"` // Current position in the movie
	Duration      string `json:"duration,omitempty"` // Movie duration
}

type FileEntry struct {
	Filename string `json:"filename"`
	IsDir    bool   `json:"directory"`
}

var (
	// Regular expression to match all supported video files
	RegexFormats = regexp.MustCompile(`.(avi|mpg|mov|flv|wmv|asf|mpeg|m4v|divx|mp4|mkv)$`)

	// Regular expression to convert filenames to titles
	RegexBrackets = regexp.MustCompile(`[\(\[\]\)]`)
	RegexYear     = regexp.MustCompile(`((19|20)[\d]{2})`)
	RegexEpisode  = regexp.MustCompile(`(?i)S[\d]+E[\d]+`)
	RegexJunk     = regexp.MustCompile(`(?i)(1080p|720p|3d|brrip|bluray|webrip|x264|aac)`)
	RegexSpace    = regexp.MustCompile(`\s{2,}`)

	// OMXPlayer control commands, these are piped via STDIN to omxplayer process
	Commands = map[string]string{
		"pause":             "p",            // Pause/continue playback
		"stop":              "q",            // Stop playback and exit
		"volume_up":         "+",            // Change volume by +3dB
		"volume_down":       "-",            // Change volume by -3dB
		"subtitles":         "s",            // Enable/disable subtitles
		"seek_back":         "\x1b\x5b\x44", // Seek -30 seconds
		"seek_back_fast":    "\x1b\x5b\x42", // Seek -600 second
		"seek_forward":      "\x1b\x5b\x43", // Seek +30 second
		"seek_forward_fast": "\x1b\x5b\x41", // Seek +600 seconds
	}

	// Path where all media files are stored
	MediaPath string

	// Path to omxplayer executable
	OmxPath string

	// Child process for spawning omxplayer
	Omx *exec.Cmd

	// Child process STDIN pipe to send commands
	OmxIn io.WriteCloser

	// Channel to pass along commands to the player routine
	Command chan string

	// Currently playing media file name
	CurrentFile string

	// Enable Zeroconf discovery
	Zeroconf bool

	// Current stream
	stream *Stream
)

// Returns true if specified file exists
func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

// Scan given path for all directories and matching video files.
// If nothing was found it will return an empty slice.
func scanPath(path string) []FileEntry {
	entries := make([]FileEntry, 0)

	files, err := ioutil.ReadDir(path)
	if err != nil {
		return entries
	}

	for _, file := range files {
		entry := FileEntry{
			Filename: file.Name(),
			IsDir:    file.IsDir(),
		}

		// Do not include non-video files in the list
		if !file.IsDir() && !omxCanPlay(file.Name()) {
			continue
		}

		entries = append(entries, entry)
	}

	return entries
}

// Convert media filename to regular title
func fileToTitle(name string) string {
	// Remove file extension from name
	name = strings.Replace(name, filepath.Ext(name), "", 1)

	// Replace all dots with white space
	name = strings.Replace(name, ".", " ", -1)

	// Replace parenteses and brackets
	name = RegexBrackets.ReplaceAllString(name, "")

	// Check if name has a typical torrent format: "name year format etc"
	pos := RegexYear.FindStringIndex(name)
	if len(pos) > 0 {
		name = name[0:pos[0]]
	} else {
		// Check if its an episode, i.e. "S##E##"
		pos = RegexEpisode.FindStringIndex(name)

		if len(pos) > 0 {
			name = name[0:pos[0]]
		}
	}

	// Remove junk stuff
	name = RegexJunk.ReplaceAllString(name, "")

	// Remove any extra white space
	name = RegexSpace.ReplaceAllString(name, "")

	// Truncate space
	name = strings.TrimSpace(name)

	return name
}

// Determine the full path to omxplayer executable. Returns error if not found.
func omxDetect() error {
	buff, err := exec.Command("which", "omxplayer").Output()
	if err != nil {
		return err
	}

	// Set path in global variable
	OmxPath = strings.TrimSpace(string(buff))

	return nil
}

// Start command listener. Commands are coming in through a channel.
func omxListen() {
	Command = make(chan string)

	for {
		command := <-Command

		// Skip command handling of omx player is not active
		if Omx == nil {
			continue
		}

		// Send command to the player
		omxWrite(command)

		// Attempt to kill the process if stop command is requested
		if command == "stop" {
			Omx.Process.Kill()
		}
	}
}

// Start omxplayer playback for a given video file. Returns error if start fails.
func omxPlay(file string) error {
	Omx = exec.Command(
		OmxPath,       // path to omxplayer executable
		"--stats",     // print stats to stdout (buffers, time, etc)
		"--with-info", // print stats about streams before playback
		"--refresh",   // adjust framerate/resolution to video
		"--blank",     // set background to black
		"--adev",      // audio out device
		"hdmi",        // using hdmi for audio/video
		file,          // path to video file
	)

	// Grab child process STDIN
	stdin, err := Omx.StdinPipe()
	if err != nil {
		return err
	}
	defer stdin.Close()

	stderr, err := Omx.StderrPipe()
	if err != nil {
		return err
	}
	defer stderr.Close()

	stdout, err := Omx.StdoutPipe()
	if err != nil {
		return err
	}
	defer stdout.Close()

	stream = NewStream()
	go stream.Start(stdout, stderr)

	// Start omxplayer execution.
	// If successful, something will appear on HDMI display.
	err = Omx.Start()
	if err != nil {
		return err
	}

	// Set current file
	CurrentFile = file

	// Make child's STDIN globally available
	OmxIn = stdin

	// Wait until child process is finished
	err = Omx.Wait()
	if err != nil {
		fmt.Fprintln(os.Stdout, "Process exited with error:", err)
	}

	omxCleanup()

	return nil
}

// Write a command string to the omxplayer process's STDIN
func omxWrite(command string) {
	if OmxIn != nil {
		io.WriteString(OmxIn, Commands[command])
	}
}

// Terminate any running omxplayer processes. Fixes random hangs.
func omxKill() {
	exec.Command("killall", "omxplayer.bin").Output()
	exec.Command("killall", "omxplayer").Output()
}

// Reset internal state and stop any running processes
func omxCleanup() {
	Omx = nil
	OmxIn = nil
	CurrentFile = ""

	omxKill()

	stream = nil
}

// Check if player is currently active
func omxIsActive() bool {
	return Omx != nil
}

// Check if player can play the file
func omxCanPlay(path string) bool {
	if RegexFormats.Match([]byte(path)) {
		return true
	}

	return false
}

// Check if TheMovieDB API could be used
func movieDbApiKey() string {
	return os.Getenv("TMDB_API_KEY")
}

func httpBrowse(c *gin.Context) {
	path := c.Request.FormValue("path")

	if path != "" {
		path = fmt.Sprintf("%s/%s", MediaPath, path)
	} else {
		path = MediaPath
	}

	c.JSON(200, scanPath(path))
}

func httpCommand(c *gin.Context) {
	val := c.Params.ByName("command")

	if _, ok := Commands[val]; !ok {
		c.JSON(400, Response{false, "Invalid command"})
		return
	}

	fmt.Println("Received command:", val)

	// Handle requested commmand
	Command <- val

	c.JSON(200, Response{true, "OK"})
}

func httpServe(c *gin.Context) {
	file := c.Request.URL.Query().Get("file")
	if file == "" {
		return
	}

	file = fmt.Sprintf("%s/%s", MediaPath, file)
	if !fileExists(file) {
		c.String(404, "Not found")
		return
	}

	if !omxCanPlay(file) {
		c.String(400, "Invalid format")
		return
	}

	http.ServeFile(c.Writer, c.Request, file)
}

func httpPlay(c *gin.Context) {
	if omxIsActive() {
		c.JSON(400, Response{false, "Player is already running"})
		return
	}

	file := c.Request.FormValue("file")
	if file == "" {
		c.JSON(400, Response{false, "File is required"})
		return
	}

	file = fmt.Sprintf("%s/%s", MediaPath, file)

	if !fileExists(file) {
		c.JSON(400, Response{false, "File does not exist"})
		return
	}

	if !omxCanPlay(file) {
		c.JSON(400, Response{false, "File cannot be played"})
		return
	}

	go omxPlay(file)

	c.JSON(200, Response{true, "OK"})
}

func httpStatus(c *gin.Context) {
	resp := StatusResponse{
		Running:       omxIsActive(),
		File:          CurrentFile,
		Name:          fileToTitle(filepath.Base(CurrentFile)),
		MovieDbApiKey: movieDbApiKey(),
	}

	if stream != nil {
		resp.Duration = durationFromSeconds(stream.duration)
		resp.Position = stream.pos.String()
	}

	c.JSON(200, resp)
}

func httpIndex(c *gin.Context) {
	data, err := Asset("static/index.html")

	if err != nil {
		c.String(400, err.Error())
		return
	}

	c.Data(200, "text/html; charset=utf-8", data)
}

func terminate(message string, code int) {
	fmt.Println(message)
	os.Exit(code)
}

func usage() {
	terminate("Usage: omxremote path/to/media/dir", 0)
}

func init() {
	var printVersion bool

	flag.StringVar(&MediaPath, "media", "./", "Path to media files")
	flag.BoolVar(&Zeroconf, "zeroconf", false, "Enable service advertisement with Zeroconf")
	flag.BoolVar(&printVersion, "v", false, "Print version")
	flag.Parse()

	if printVersion {
		fmt.Printf("omxremote v%v\n", VERSION)
		os.Exit(0)
	}
}

func main() {
	// Expand media path if needed
	MediaPath = strings.Replace(MediaPath, "~", os.Getenv("HOME"), 1)

	// Get path from arguments and remove trailing slash
	MediaPath = strings.TrimRight(MediaPath, "/")

	if !fileExists(MediaPath) {
		terminate(fmt.Sprintf("Directory does not exist: %s", MediaPath), 1)
	}

	// Check if player is installed
	if omxDetect() != nil {
		terminate("omxplayer is not installed", 1)
	}

	// Make sure nothing is running
	omxCleanup()

	// Start a remote command listener
	go omxListen()

	// Start zeroconf service advertisement
	if (Zeroconf) {
		stopZeroconf := make(chan bool)
		go startZeroConfAdvertisement(stopZeroconf)
	}

	// Disable debugging mode
	gin.SetMode("release")

	// Setup HTTP server
	router := gin.Default()

	// Handle CORS
	router.Use(func(c *gin.Context) {
		c.Header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
		c.Header("Access-Control-Allow-Origin", "*")
		c.Header("Access-Control-Expose-Headers", "*")
	})

	router.GET("/", httpIndex)
	router.GET("/status", httpStatus)
	router.GET("/browse", httpBrowse)
	router.GET("/play", httpPlay)
	router.GET("/serve", httpServe)
	router.GET("/command/:command", httpCommand)

	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}

	fmt.Println("Starting server on 0.0.0.0:" + port)
	router.Run(":" + port)
}
<|endoftext|>


<commit_before>package negroni

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"runtime/debug"
)

type ErrorHandlerFunc interface {
	Handle(err interface{}, report bool)
}

// Recovery is a Negroni middleware that recovers from any panics and writes a 500 if there was one.
type Recovery struct {
	Logger           *log.Logger
	PrintStack       bool
	ErrorHandlerFunc ErrorHandlerFunc
}

// NewRecovery returns a new instance of Recovery
func NewRecovery() *Recovery {
	return &Recovery{
		Logger:     log.New(os.Stdout, "[negroni] ", 0),
		PrintStack: true,
	}
}

func (rec *Recovery) ServeHTTP(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	defer func() {
		if err := recover(); err != nil {
			rw.WriteHeader(http.StatusInternalServerError)
			stack := debug.Stack()
			f := "PANIC: %s\n%s"
			rec.Logger.Printf(f, err, stack)

			if rec.PrintStack {
				fmt.Fprintf(rw, f, err, stack)
			}

			if rec.ErrorHandlerFunc != nil {
				rec.ErrorHandlerFunc.Handle(err, true)
			}
		}
	}()

	next(rw, r)
}
<commit_msg>Add Optional ErrorHandlerFunc<commit_after>package negroni

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"runtime/debug"
)

type ErrorHandlerFunc interface {
	Handle(err interface{}, report bool)
}

// Recovery is a Negroni middleware that recovers from any panics and writes a 500 if there was one.
type Recovery struct {
	Logger           *log.Logger
	PrintStack       bool
	ErrorHandlerFunc ErrorHandlerFunc
}

// NewRecovery returns a new instance of Recovery
func NewRecovery() *Recovery {
	return &Recovery{
		Logger:     log.New(os.Stdout, "[negroni] ", 0),
		PrintStack: true,
	}
}

// NewRecoveryWithFunc returns a nwe instance of Recovery with an
// attached errorHandlerFunction
func NewRecoveryWithFunc(errorHandlerFunc ErrorHandlerFunc) *Recovery {
	return &Recovery{
		Logger:           log.New(os.Stdout, "[negroni] ", 0),
		PrintStack:       true,
		ErrorHandlerFunc: errorHandlerFunc,
	}
}

func (rec *Recovery) ServeHTTP(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	defer func() {
		if err := recover(); err != nil {
			rw.WriteHeader(http.StatusInternalServerError)
			stack := debug.Stack()
			f := "PANIC: %s\n%s"
			rec.Logger.Printf(f, err, stack)

			if rec.PrintStack {
				fmt.Fprintf(rw, f, err, stack)
			}

			if rec.ErrorHandlerFunc != nil {
				rec.ErrorHandlerFunc.Handle(err, true)
			}
		}
	}()

	next(rw, r)
}
<|endoftext|>


<commit_before>package mirango

import (
	"fmt"

	"github.com/mirango/defaults"
	"github.com/mirango/framework"
)

type Operations struct {
	operations []*Operation
	index      map[string]int
}

func NewOperations() *Operations {
	return &Operations{
		index: map[string]int{},
	}
}

func (ops *Operations) Append(operations ...*Operation) {
	le := len(ops.operations)
	for i := 0; i < len(operations); i++ {
		name := operations[i].name
		if name == "" {
			ops.operations = append(ops.operations, operations[i])
			continue
		} else {
			if _, ok := ops.index[name]; ok {
				panic(fmt.Sprintf("Detected 2 operations with the same name: \"%s\".", name))
			}
			ops.operations = append(ops.operations, operations[i])
			ops.index[name] = le + i
		}
	}
}

func (ops *Operations) Set(operations ...*Operation) {
	ops.operations = nil
	ops.index = map[string]int{}
	ops.Append(operations...)
}

func (ops *Operations) Get(name string) *Operation {
	return ops.operations[ops.index[name]]
}

func (ops *Operations) GetAll() []*Operation {
	return ops.operations
}

func (ops *Operations) Union(nops *Operations) {
	for _, o := range nops.operations {
		ops.Append(o)
	}
}

func (ops *Operations) Clone() *Operations {
	nops := NewOperations()
	nops.Union(ops)
	return nops
}

func (ops *Operations) GetByIndex(i int) *Operation {
	return ops.operations[i]
}

func (ops *Operations) GetByMethod(method string) *Operation {
	for _, o := range ops.operations {
		for _, m := range o.methods {
			if m == method {
				return o
			}
		}
	}
	return nil
}

func (ops *Operations) Len() int {
	return len(ops.operations)
}

type Operation struct {
	name          string
	handler       Handler
	route         *Route
	methods       []string
	schemes       []string
	accepts       []string
	returns       []string
	middleware    []Middleware
	params        *Params
	renders       string
	mimeTypeIn    paramIn
	mimeTypeParam string

	allParams  *Params
	allSchemes []string
	allAccepts []string
	allReturns []string
}

func NewOperation(r *Route, h interface{}) *Operation {
	handler, err := handler(h)
	if err != nil {
		panic(err.Error())
	}
	o := &Operation{
		methods: []string{"GET"},
		handler: handler,
		route:   r,
		params:  NewParams(),
		schemes: defaults.Schemes,
		accepts: defaults.Accepts,
		returns: defaults.Returns,
	}
	o.middleware = []Middleware{CheckReturns(o), CheckSchemes(o), CheckAccepts(o), CheckParams(o)}
	return o
}

func GET(r *Route, h interface{}) *Operation {
	return NewOperation(r, h).Methods("GET")
}

func POST(r *Route, h interface{}) *Operation {
	return NewOperation(r, h).Methods("POST")
}

func PUT(r *Route, h interface{}) *Operation {
	return NewOperation(r, h).Methods("PUT")
}

func DELETE(r *Route, h interface{}) *Operation {
	return NewOperation(r, h).Methods("DELETE")
}

func (o *Operation) Uses(h interface{}) *Operation { //interface
	handler, err := handler(h)
	if err != nil {
		panic(err.Error())
	}
	o.handler = handler
	return o
}

func getHandler(h interface{}, mw []Middleware) (Handler, error) {
	final, err := handler(h)
	if err != nil {
		return nil, err
	}
	for i := len(mw) - 1; i >= 0; i-- {
		final = mw[i].Run(final)
	}
	return final, nil
}

func (o *Operation) With(mw ...interface{}) *Operation {
	for i := 0; i < len(mw); i++ {
		switch t := mw[i].(type) {
		case Middleware:
			o.middleware = append(o.middleware, t)
		case MiddlewareFunc:
			o.middleware = append(o.middleware, t)
		case func(Handler) Handler:
			o.middleware = append(o.middleware, MiddlewareFunc(t))
		}
	}
	return o
}

func (o *Operation) with() {
	handler, err := getHandler(o.handler, o.middleware)
	if err != nil {
		panic(err.Error())
	}
	o.handler = handler
}

func (o *Operation) Apply(temps ...func(*Operation)) *Operation {
	for i := 0; i < len(temps); i++ {
		temps[i](o)
	}
	return o
}

func (o *Operation) Route() framework.Route {
	return o.route
}

func (o *Operation) Name(name string) *Operation {
	o.name = name
	return o
}

func (o *Operation) GetName() string {
	return o.name
}

func (o *Operation) Methods(methods ...string) *Operation {
	o.methods = methods
	return o
}

func (o *Operation) Params(params ...*Param) *Operation {
	o.params.Set(params...)
	return o
}

func (o *Operation) GetParams() *Params {
	return o.params
}

func (o *Operation) GetAllParams() *Params {
	if o.allParams != nil {
		return o.allParams
	}
	params := o.params.Clone()
	params.Union(o.route.GetAllParams())
	o.allParams = params
	o.params = nil
	return params
}

func (o *Operation) Schemes(schemes ...string) *Operation {
	o.schemes = schemes
	return o
}

func (o *Operation) Accepts(accepts ...string) *Operation {
	o.accepts = accepts
	return o
}

func (o *Operation) Returns(returns ...string) *Operation {
	o.returns = returns
	return o
}

func (o *Operation) PathParam(name string) *Param {
	p := PathParam(name)
	o.params.Append(p)
	return p
}

func (o *Operation) QueryParam(name string) *Param {
	p := QueryParam(name)
	o.params.Append(p)
	return p
}

func (o *Operation) HeaderParam(name string) *Param {
	p := HeaderParam(name)
	o.params.Append(p)
	return p
}

func (o *Operation) BodyParam(name string) *Param {
	p := BodyParam(name)
	o.params.Append(p)
	return p
}

func (o *Operation) GetMethods() []string {
	return o.methods
}

func (o *Operation) GetMiddleware() []Middleware {
	return o.middleware
}

func (o *Operation) GetAllMiddleware() []Middleware {
	return middlewareUnion(o.middleware, o.route.GetAllMiddleware())
}

func (o *Operation) GetSchemes() []string {
	return o.methods
}

func (o *Operation) GetAllSchemes() []string {
	if o.allSchemes != nil {
		return o.allSchemes
	}
	schemes := stringsUnion(o.schemes, o.route.GetAllSchemes())
	o.allSchemes = schemes
	o.schemes = nil
	return schemes
}

func (o *Operation) GetAccepts() []string {
	return o.methods
}

func (o *Operation) GetAllAccepts() []string {
	if o.allAccepts != nil {
		return o.allAccepts
	}
	accepts := stringsUnion(o.accepts, o.route.GetAllAccepts())
	o.allAccepts = accepts
	o.accepts = nil
	return accepts
}

func (o *Operation) GetReturns() []string {
	return o.methods
}

func (o *Operation) GetAllReturns() []string {
	if o.allReturns != nil {
		return o.allReturns
	}
	returns := stringsUnion(o.returns, o.route.GetAllReturns())
	o.allReturns = returns
	o.returns = nil
	return returns
}

func (o *Operation) BuildPath(v ...interface{}) string {
	return o.route.BuildPath(v...)
}

func (o *Operation) GetPath() string {
	return o.route.path
}

func (o *Operation) GetFullPath() string {
	return o.route.GetFullPath()
}

func (o *Operation) ServeHTTP(c *Context) interface{} {
	if o.middleware != nil {
		o.middleware = o.GetAllMiddleware()
		o.with()
		o.middleware = nil
	}
	c.operation = o
	return o.handler.ServeHTTP(c)
}

func (o *Operation) Clone() *Operation {
	no := NewOperation(nil, o.handler)

	no.methods = o.methods
	no.schemes = o.schemes
	no.accepts = o.accepts
	no.returns = o.returns
	no.middleware = o.middleware
	no.params = o.params.Clone()
	no.renders = o.renders
	no.mimeTypeIn = o.mimeTypeIn
	no.mimeTypeParam = o.mimeTypeParam

	return no
}

type middlewareContainer struct {
	middleware []interface{}
}

func With(mw ...interface{}) *middlewareContainer {
	return &middlewareContainer{middleware: mw}
}

func (c middlewareContainer) Handle(operations ...*Operation) []*Operation {
	for i := 0; i < len(operations); i++ {
		operations[i].With(c.middleware...)
	}
	return operations
}

type templateContainer struct {
	templates []func(*Operation)
}

func Apply(temps ...func(*Operation)) *templateContainer {
	return &templateContainer{templates: temps}
}

func (c templateContainer) To(operations ...*Operation) []*Operation {
	for i := 0; i < len(operations); i++ {
		operations[i].Apply(c.templates...)
	}
	return operations
}
<commit_msg>Added ReturnsOnly, AcceptsOnly and SchemesOnly methods<commit_after>package mirango

import (
	"fmt"

	"github.com/mirango/defaults"
	"github.com/mirango/framework"
)

type Operations struct {
	operations []*Operation
	index      map[string]int
}

func NewOperations() *Operations {
	return &Operations{
		index: map[string]int{},
	}
}

func (ops *Operations) Append(operations ...*Operation) {
	le := len(ops.operations)
	for i := 0; i < len(operations); i++ {
		name := operations[i].name
		if name == "" {
			ops.operations = append(ops.operations, operations[i])
			continue
		} else {
			if _, ok := ops.index[name]; ok {
				panic(fmt.Sprintf("Detected 2 operations with the same name: \"%s\".", name))
			}
			ops.operations = append(ops.operations, operations[i])
			ops.index[name] = le + i
		}
	}
}

func (ops *Operations) Set(operations ...*Operation) {
	ops.operations = nil
	ops.index = map[string]int{}
	ops.Append(operations...)
}

func (ops *Operations) Get(name string) *Operation {
	return ops.operations[ops.index[name]]
}

func (ops *Operations) GetAll() []*Operation {
	return ops.operations
}

func (ops *Operations) Union(nops *Operations) {
	for _, o := range nops.operations {
		ops.Append(o)
	}
}

func (ops *Operations) Clone() *Operations {
	nops := NewOperations()
	nops.Union(ops)
	return nops
}

func (ops *Operations) GetByIndex(i int) *Operation {
	return ops.operations[i]
}

func (ops *Operations) GetByMethod(method string) *Operation {
	for _, o := range ops.operations {
		for _, m := range o.methods {
			if m == method {
				return o
			}
		}
	}
	return nil
}

func (ops *Operations) Len() int {
	return len(ops.operations)
}

type Operation struct {
	name          string
	handler       Handler
	route         *Route
	methods       []string
	schemes       []string
	accepts       []string
	returns       []string
	middleware    []Middleware
	params        *Params
	renders       string
	mimeTypeIn    paramIn
	mimeTypeParam string

	allParams  *Params
	allSchemes []string
	allAccepts []string
	allReturns []string

	returnsOnly bool
	acceptsOnly bool
	schemesOnly bool
}

func NewOperation(r *Route, h interface{}) *Operation {
	handler, err := handler(h)
	if err != nil {
		panic(err.Error())
	}
	o := &Operation{
		methods: []string{"GET"},
		handler: handler,
		route:   r,
		params:  NewParams(),
		schemes: defaults.Schemes,
		accepts: defaults.Accepts,
		returns: defaults.Returns,
	}
	o.middleware = []Middleware{CheckReturns(o), CheckSchemes(o), CheckAccepts(o), CheckParams(o)}
	return o
}

func GET(r *Route, h interface{}) *Operation {
	return NewOperation(r, h).Methods("GET")
}

func POST(r *Route, h interface{}) *Operation {
	return NewOperation(r, h).Methods("POST")
}

func PUT(r *Route, h interface{}) *Operation {
	return NewOperation(r, h).Methods("PUT")
}

func DELETE(r *Route, h interface{}) *Operation {
	return NewOperation(r, h).Methods("DELETE")
}

func (o *Operation) Uses(h interface{}) *Operation { //interface
	handler, err := handler(h)
	if err != nil {
		panic(err.Error())
	}
	o.handler = handler
	return o
}

func getHandler(h interface{}, mw []Middleware) (Handler, error) {
	final, err := handler(h)
	if err != nil {
		return nil, err
	}
	for i := len(mw) - 1; i >= 0; i-- {
		final = mw[i].Run(final)
	}
	return final, nil
}

func (o *Operation) With(mw ...interface{}) *Operation {
	for i := 0; i < len(mw); i++ {
		switch t := mw[i].(type) {
		case Middleware:
			o.middleware = append(o.middleware, t)
		case MiddlewareFunc:
			o.middleware = append(o.middleware, t)
		case func(Handler) Handler:
			o.middleware = append(o.middleware, MiddlewareFunc(t))
		}
	}
	return o
}

func (o *Operation) with() {
	handler, err := getHandler(o.handler, o.middleware)
	if err != nil {
		panic(err.Error())
	}
	o.handler = handler
}

func (o *Operation) Apply(temps ...func(*Operation)) *Operation {
	for i := 0; i < len(temps); i++ {
		temps[i](o)
	}
	return o
}

func (o *Operation) Route() framework.Route {
	return o.route
}

func (o *Operation) Name(name string) *Operation {
	o.name = name
	return o
}

func (o *Operation) GetName() string {
	return o.name
}

func (o *Operation) Methods(methods ...string) *Operation {
	o.methods = methods
	return o
}

func (o *Operation) Params(params ...*Param) *Operation {
	o.params.Set(params...) // make it append instead of set
	return o
}

func (o *Operation) GetParams() *Params {
	return o.params
}

func (o *Operation) GetAllParams() *Params {
	if o.allParams != nil {
		return o.allParams
	}
	params := o.params.Clone()
	params.Union(o.route.GetAllParams())
	o.allParams = params
	o.params = nil
	return params
}

func (o *Operation) Schemes(schemes ...string) *Operation {
	o.schemes = append(o.schemes, schemes...)
	o.schemesOnly = false
	return o
}

func (o *Operation) Accepts(accepts ...string) *Operation {
	o.accepts = append(o.accepts, accepts...)
	o.acceptsOnly = false
	return o
}

func (o *Operation) Returns(returns ...string) *Operation {
	o.returns = append(o.returns, returns...)
	o.returnsOnly = false
	return o
}

func (o *Operation) SchemesOnly(schemes ...string) *Operation {
	o.schemes = schemes
	o.schemesOnly = true
	return o
}

func (o *Operation) AcceptsOnly(accepts ...string) *Operation {
	o.accepts = accepts
	o.acceptsOnly = true
	return o
}

func (o *Operation) ReturnsOnly(returns ...string) *Operation {
	o.returns = returns
	o.returnsOnly = true
	return o
}

func (o *Operation) PathParam(name string) *Param {
	p := PathParam(name)
	o.params.Append(p)
	return p
}

func (o *Operation) QueryParam(name string) *Param {
	p := QueryParam(name)
	o.params.Append(p)
	return p
}

func (o *Operation) HeaderParam(name string) *Param {
	p := HeaderParam(name)
	o.params.Append(p)
	return p
}

func (o *Operation) BodyParam(name string) *Param {
	p := BodyParam(name)
	o.params.Append(p)
	return p
}

func (o *Operation) GetMethods() []string {
	return o.methods
}

func (o *Operation) GetMiddleware() []Middleware {
	return o.middleware
}

func (o *Operation) GetAllMiddleware() []Middleware {
	return middlewareUnion(o.middleware, o.route.GetAllMiddleware())
}

func (o *Operation) GetSchemes() []string {
	return o.methods
}

func (o *Operation) GetAllSchemes() []string {
	if o.allSchemes != nil {
		return o.allSchemes
	}
	schemes := o.schemes
	if !o.schemesOnly {
		schemes = stringsUnion(schemes, o.route.GetAllSchemes())
	}
	o.allSchemes = schemes
	o.schemes = nil
	return schemes
}

func (o *Operation) GetAccepts() []string {
	return o.methods
}

func (o *Operation) GetAllAccepts() []string {
	if o.allAccepts != nil {
		return o.allAccepts
	}
	accepts := o.accepts
	if !o.acceptsOnly {
		accepts = stringsUnion(accepts, o.route.GetAllAccepts())
	}
	o.allAccepts = accepts
	o.accepts = nil
	return accepts
}

func (o *Operation) GetReturns() []string {
	return o.methods
}

func (o *Operation) GetAllReturns() []string {
	if o.allReturns != nil {
		return o.allReturns
	}
	returns := o.returns
	if !o.returnsOnly {
		returns = stringsUnion(returns, o.route.GetAllReturns())
	}
	o.allReturns = returns
	o.returns = nil
	return returns
}

func (o *Operation) BuildPath(v ...interface{}) string {
	return o.route.BuildPath(v...)
}

func (o *Operation) GetPath() string {
	return o.route.path
}

func (o *Operation) GetFullPath() string {
	return o.route.GetFullPath()
}

func (o *Operation) ServeHTTP(c *Context) interface{} {
	c.operation = o
	if o.middleware != nil {
		o.middleware = o.GetAllMiddleware()
		o.with()
		o.middleware = nil
	}
	return o.handler.ServeHTTP(c)
}

func (o *Operation) Clone() *Operation {
	no := NewOperation(nil, o.handler)

	no.methods = o.methods
	no.schemes = o.schemes
	no.accepts = o.accepts
	no.returns = o.returns
	no.middleware = o.middleware
	no.params = o.params.Clone()
	no.renders = o.renders
	no.mimeTypeIn = o.mimeTypeIn
	no.mimeTypeParam = o.mimeTypeParam

	return no
}

type middlewareContainer struct {
	middleware []interface{}
}

func With(mw ...interface{}) *middlewareContainer {
	return &middlewareContainer{middleware: mw}
}

func (c middlewareContainer) Handle(operations ...*Operation) []*Operation {
	for i := 0; i < len(operations); i++ {
		operations[i].With(c.middleware...)
	}
	return operations
}

type templateContainer struct {
	templates []func(*Operation)
}

func Apply(temps ...func(*Operation)) *templateContainer {
	return &templateContainer{templates: temps}
}

func (c templateContainer) To(operations ...*Operation) []*Operation {
	for i := 0; i < len(operations); i++ {
		operations[i].Apply(c.templates...)
	}
	return operations
}
<|endoftext|>