<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"
)

type testStatus struct {
	Success bool   `json:"success"`
	Latency int64  `json:"latency"`
	Error   string `json:"error"`
}

type toxStatus struct {
	LastScan    int64      `json:"last_scan"`
	LastRefresh int64      `json:"last_refresh"`
	Nodes       []*toxNode `json:"nodes"`
}

type toxNode struct {
	Ipv4Address string `json:"ipv4"`
	Ipv6Address string `json:"ipv6"`
	Port        int    `json:"port"`
	TCPPorts    []int  `json:"tcp_ports"`
	PublicKey   string `json:"public_key"`
	Maintainer  string `json:"maintainer"`
	Location    string `json:"location"`
	UDPStatus   bool   `json:"status_udp"`
	TCPStatus   bool   `json:"status_tcp"`
	Version     string `json:"version"`
	MOTD        string `json:"motd"`
	LastPing    int64  `json:"last_ping"`
	index       int
	ip4         net.IP
	ip6         net.IP
}

const (
	httpListenPort = 8081
	wikiURI        = "https://wiki.tox.chat/users/nodes?do=export_raw"
)

var (
	countries map[string]string
)

func loadCountries() error {
	const name = "countries.json"
	bytes, ok := assetMap[name]
	if !ok {
		return fmt.Errorf("asset %s not found", name)
	}

	return json.Unmarshal(bytes, &countries)
}

func handleHTTPRequest(w http.ResponseWriter, r *http.Request) {
	urlPath := r.URL.Path[1:]
	if r.URL.Path == "/" {
		urlPath = "index"
	}

	if r.Method == "POST" && r.URL.Path == "/test" {
		handleTestRequest(w, r)
		return
	}

	var content interface{}
	var filename string
	switch r.URL.Path {
	case "/":
		content = toxStatus{lastScan, lastRefresh, nodes}
		fallthrough
	case "/test", "/about":
		filename = urlPath + ".html"
	default:
		data, ok := assetMap[urlPath]
		if !ok {
			http.Error(w, http.StatusText(404), 404)
		} else {
			w.Header().Set("Content-Type", mimeTypeByExtension(urlPath))
			w.Write(data)
		}
		return
	}

	err := renderTemplate(w, filename, content)
	if err != nil {
		fmt.Printf("tmpl exec error: %s\n", err.Error())
	}
}

func handleTestRequest(w http.ResponseWriter, r *http.Request) {
	fmt.Printf("test request from: %s\n", r.RemoteAddr)

	err := r.ParseForm()
	if err != nil {
		http.Error(w, "unable to parse the form", 500)
		return
	}

	ipPort := r.Form.Get("ipPort")
	key := r.Form.Get("key")
	net := r.Form.Get("net")

	ipPortParts := strings.Split(ipPort, ":")
	if len(ipPortParts) != 2 {
		http.Error(w, "invalid ip:port submission", 500)
		return
	}

	port, err := strconv.Atoi(ipPortParts[1])
	if err != nil {
		http.Error(w, "specified port is not an int", 500)
		return
	}

	ip := ipPortParts[0]
	ip4, ip6 := resolveIPAddr(ip, ip)
	node := toxNode{
		PublicKey: key,
		Port:      port,
	}
	if ip4 != nil {
		node.ip4 = ip4.IP
	}
	if ip6 != nil {
		node.ip6 = ip6.IP
	}

	err = nil
	start := time.Now()
	content := testStatus{Success: true}

	switch net {
	case "UDP":
		err = probeNode(&node)
	case "TCP":
		err = probeNodeTCP(&node)
	default:
		http.Error(w, "invalid net type", 500)
		return
	}

	if err != nil {
		content.Success = false
		content.Error = err.Error()
	} else {
		content.Latency = time.Now().Sub(start).Nanoseconds() / int64(time.Millisecond)
	}

	writeJSONResponse(w, content)
}

func handleJSONRequest(w http.ResponseWriter, r *http.Request) {
	content := toxStatus{lastScan, lastRefresh, nodes}
	writeJSONResponse(w, content)
}

func writeJSONResponse(w http.ResponseWriter, content interface{}) {
	bytes, err := json.Marshal(content)
	if err != nil {
		http.Error(w, http.StatusText(500), 500)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(bytes)
}

func parseNodes() ([]*toxNode, error) {
	client := http.Client{}
	client.Timeout = 5 * time.Second
	res, err := client.Get(wikiURI)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	nodesList := []*toxNode{}
	content, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	lines := strings.Split(string(content), "\n")
	for i, line := range lines {
		node := parseNode(line)
		if node == nil {
			continue
		}

		node.index = i
		nodesList = append(nodesList, node)
	}
	return nodesList, nil
}

func parseNode(nodeString string) *toxNode {
	nodeString = stripSpaces(nodeString)
	if !strings.HasPrefix(nodeString, "|") {
		return nil
	}

	lineParts := strings.Split(nodeString, "|")
	var node *toxNode

	if port, err := strconv.Atoi(lineParts[3]); err == nil && len(lineParts) == 8 {
		ip4, ip6 := resolveIPAddr(lineParts[1], lineParts[2])

		node = &toxNode{
			Ipv4Address: lineParts[1],
			Ipv6Address: lineParts[2],
			Port:        port,
			TCPPorts:    []int{},
			PublicKey:   lineParts[4],
			Maintainer:  lineParts[5],
			Location:    lineParts[6],
			UDPStatus:   false,
			TCPStatus:   false,
			Version:     "",
			MOTD:        "",
			LastPing:    0,
			ip4:         nil,
			ip6:         nil,
		}

		if ip4 != nil {
			node.ip4 = ip4.IP
		}
		if ip6 != nil {
			node.ip6 = ip6.IP
		}

		if node.Ipv6Address == "NONE" {
			node.Ipv6Address = "-"
		}
	}

	return node
}

func resolveIPAddr(ip4String string, ip6String string) (*net.IPAddr, *net.IPAddr) {
	ip4, err := net.ResolveIPAddr("ip4", ip4String)
	if err != nil {
		fmt.Printf("couldn't resolve %s: %s\n", ip4String, err.Error())
	}

	ip6, err := net.ResolveIPAddr("ip6", ip6String)
	if err != nil {
		fmt.Printf("couldn't resolve %s: %s\n", ip6String, err.Error())
	}

	return ip4, ip6
}
<commit_msg>Don't bother trying to resolve if a node's ip is "NONE"<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"
)

type testStatus struct {
	Success bool   `json:"success"`
	Latency int64  `json:"latency"`
	Error   string `json:"error"`
}

type toxStatus struct {
	LastScan    int64      `json:"last_scan"`
	LastRefresh int64      `json:"last_refresh"`
	Nodes       []*toxNode `json:"nodes"`
}

type toxNode struct {
	Ipv4Address string `json:"ipv4"`
	Ipv6Address string `json:"ipv6"`
	Port        int    `json:"port"`
	TCPPorts    []int  `json:"tcp_ports"`
	PublicKey   string `json:"public_key"`
	Maintainer  string `json:"maintainer"`
	Location    string `json:"location"`
	UDPStatus   bool   `json:"status_udp"`
	TCPStatus   bool   `json:"status_tcp"`
	Version     string `json:"version"`
	MOTD        string `json:"motd"`
	LastPing    int64  `json:"last_ping"`
	index       int
	ip4         net.IP
	ip6         net.IP
}

const (
	httpListenPort = 8081
	wikiURI        = "https://wiki.tox.chat/users/nodes?do=export_raw"
)

var (
	countries map[string]string
)

func loadCountries() error {
	const name = "countries.json"
	bytes, ok := assetMap[name]
	if !ok {
		return fmt.Errorf("asset %s not found", name)
	}

	return json.Unmarshal(bytes, &countries)
}

func handleHTTPRequest(w http.ResponseWriter, r *http.Request) {
	urlPath := r.URL.Path[1:]
	if r.URL.Path == "/" {
		urlPath = "index"
	}

	if r.Method == "POST" && r.URL.Path == "/test" {
		handleTestRequest(w, r)
		return
	}

	var content interface{}
	var filename string
	switch r.URL.Path {
	case "/":
		content = toxStatus{lastScan, lastRefresh, nodes}
		fallthrough
	case "/test", "/about":
		filename = urlPath + ".html"
	default:
		data, ok := assetMap[urlPath]
		if !ok {
			http.Error(w, http.StatusText(404), 404)
		} else {
			w.Header().Set("Content-Type", mimeTypeByExtension(urlPath))
			w.Write(data)
		}
		return
	}

	err := renderTemplate(w, filename, content)
	if err != nil {
		fmt.Printf("tmpl exec error: %s\n", err.Error())
	}
}

func handleTestRequest(w http.ResponseWriter, r *http.Request) {
	fmt.Printf("test request from: %s\n", r.RemoteAddr)

	err := r.ParseForm()
	if err != nil {
		http.Error(w, "unable to parse the form", 500)
		return
	}

	ipPort := r.Form.Get("ipPort")
	key := r.Form.Get("key")
	net := r.Form.Get("net")

	ipPortParts := strings.Split(ipPort, ":")
	if len(ipPortParts) != 2 {
		http.Error(w, "invalid ip:port submission", 500)
		return
	}

	port, err := strconv.Atoi(ipPortParts[1])
	if err != nil {
		http.Error(w, "specified port is not an int", 500)
		return
	}

	ip := ipPortParts[0]
	ip4, ip6 := resolveIPAddr(ip, ip)
	node := toxNode{
		PublicKey: key,
		Port:      port,
	}
	if ip4 != nil {
		node.ip4 = ip4.IP
	}
	if ip6 != nil {
		node.ip6 = ip6.IP
	}

	err = nil
	start := time.Now()
	content := testStatus{Success: true}

	switch net {
	case "UDP":
		err = probeNode(&node)
	case "TCP":
		err = probeNodeTCP(&node)
	default:
		http.Error(w, "invalid net type", 500)
		return
	}

	if err != nil {
		content.Success = false
		content.Error = err.Error()
	} else {
		content.Latency = time.Now().Sub(start).Nanoseconds() / int64(time.Millisecond)
	}

	writeJSONResponse(w, content)
}

func handleJSONRequest(w http.ResponseWriter, r *http.Request) {
	content := toxStatus{lastScan, lastRefresh, nodes}
	writeJSONResponse(w, content)
}

func writeJSONResponse(w http.ResponseWriter, content interface{}) {
	bytes, err := json.Marshal(content)
	if err != nil {
		http.Error(w, http.StatusText(500), 500)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(bytes)
}

func parseNodes() ([]*toxNode, error) {
	client := http.Client{}
	client.Timeout = 5 * time.Second
	res, err := client.Get(wikiURI)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	nodesList := []*toxNode{}
	content, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	lines := strings.Split(string(content), "\n")
	for i, line := range lines {
		node := parseNode(line)
		if node == nil {
			continue
		}

		node.index = i
		nodesList = append(nodesList, node)
	}
	return nodesList, nil
}

func parseNode(nodeString string) *toxNode {
	nodeString = stripSpaces(nodeString)
	if !strings.HasPrefix(nodeString, "|") {
		return nil
	}

	lineParts := strings.Split(nodeString, "|")
	var node *toxNode

	if port, err := strconv.Atoi(lineParts[3]); err == nil && len(lineParts) == 8 {
		ip4, ip6 := resolveIPAddr(lineParts[1], lineParts[2])

		node = &toxNode{
			Ipv4Address: lineParts[1],
			Ipv6Address: lineParts[2],
			Port:        port,
			TCPPorts:    []int{},
			PublicKey:   lineParts[4],
			Maintainer:  lineParts[5],
			Location:    lineParts[6],
			UDPStatus:   false,
			TCPStatus:   false,
			Version:     "",
			MOTD:        "",
			LastPing:    0,
			ip4:         nil,
			ip6:         nil,
		}

		if ip4 != nil {
			node.ip4 = ip4.IP
		}
		if ip6 != nil {
			node.ip6 = ip6.IP
		}

		if node.Ipv6Address == "NONE" {
			node.Ipv6Address = "-"
		}
	}

	return node
}

func resolveIPAddr(ip4String string, ip6String string) (*net.IPAddr, *net.IPAddr) {
	resolve := func(network string, ipString string) *net.IPAddr {
		if ipString != "NONE" {
			ip, err := net.ResolveIPAddr(network, ipString)
			if err != nil {
				fmt.Printf("couldn't resolve %s: %s\n", ip4String, err.Error())
			}
			return ip
		}
		return nil
	}

	return resolve("ip4", ip4String), resolve("ip6", ip6String)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"text/template"
)

var (
	output   string
	pkg      string
	function string
	fs       string
	doFormat bool

	tmpl = template.Must(template.New("go-template").Parse(goTemplate))
)

type templateData struct {
	Pkg      string
	Function string
	FS       string

	Files []File
}

type File struct {
	Path    string
	Data    string
	Size    int64
	Mode    uint32
	ModTime int64
}

func main() {
	flag.Parse()
	if flag.NArg() == 0 {
		fmt.Fprintln(os.Stderr, "No input files/directories given")
		os.Exit(0)
	}

	if pkg == "" {
		fmt.Fprintln(os.Stderr, "No output file package given")
		os.Exit(0)
	}

	var out *os.File

	if output == "-" {
		out = os.Stdout
	} else {
		var err error
		if out, err = os.OpenFile(output, os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644); err != nil {
			fmt.Fprintf(os.Stderr, "Error opening '%s' for writing: %v\n", output, err)
			os.Exit(0)
		}

		defer out.Close()
	}

	files := []File{}

	for _, arg := range flag.Args() {
		recursive := false
		if strings.HasSuffix(arg, "/...") {
			recursive = true
			arg = arg[:len(arg)-4]
		}
		func() {
			file, err := os.Open(arg)
			defer file.Close()

			if err != nil {
				fmt.Fprintf(os.Stderr, "Error opening '%s': %v\n", arg, err)
				return
			}

			stat, err := file.Stat()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error stat-ing '%s': %v\n", arg, err)
				return
			}

			if stat.IsDir() {
				filepath.Walk(arg, func(path string, info os.FileInfo, err error) error {
					if info.IsDir() {
						if !recursive {
							return filepath.SkipDir
						}
					} else {
						if b, err := ioutil.ReadFile(path); err == nil {
							files = append(files, File{
								Path:    path,
								Data:    fmt.Sprintf("%q", b),
								Size:    info.Size(),
								ModTime: info.ModTime().Unix(),
								Mode:    uint32(info.Mode()),
							})
						} else {
							fmt.Fprintf(os.Stderr, "Error reading '%s': %v\n", path, err)
						}
					}

					return nil
				})
			} else {
				if b, err := ioutil.ReadAll(file); err == nil {
					files = append(files, File{
						Path:    arg,
						Data:    fmt.Sprintf("%q", b),
						Size:    stat.Size(),
						ModTime: stat.ModTime().Unix(),
						Mode:    uint32(stat.Mode()),
					})
				} else {
					fmt.Fprintf(os.Stderr, "Error reading '%s': %v\n", arg, err)
					return
				}

			}
		}()
	}

	buf := new(bytes.Buffer)
	err := tmpl.Execute(buf, templateData{Pkg: pkg, Function: function, FS: fs, Files: files})
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error executing template: %v\n", err)
		os.Exit(0)
	}

	if doFormat {
		if b, err := format.Source(buf.Bytes()); err == nil {
			buf.Reset()
			if _, err = buf.Write(b); err != nil {
				fmt.Fprintf(os.Stderr, "Error writing formatted code to buffer: %v\n", err)
				os.Exit(0)
			}
		} else {
			fmt.Fprintf(os.Stderr, "Error formatting generated code: %v\n", err)
			os.Exit(0)
		}
	}

	if _, err := buf.WriteTo(out); err != nil {
		fmt.Fprintf(os.Stderr, "Error writing template result to output: %v\n", err)
		os.Exit(0)
	}
}

func init() {
	flag.StringVar(&output, "output", "-", "the output file")
	flag.StringVar(&pkg, "package", "main", "the package of the output go file")
	flag.StringVar(&function, "function", "addFiles", "the function that will add the files to the fs")
	flag.StringVar(&fs, "fs", "DefaultFS", "the fs object to be used. Will be created if different from the default")
	flag.BoolVar(&doFormat, "format", false, "run the output go code through go/format")
}

const goTemplate = `package {{ .Pkg }}

// DO NOT EDIT ** This file was generated with the webfw-fs tool ** DO NOT EDIT //

import (
	"fmt"
	"os"
	"time"

	"github.com/urandom/webfw/fs"
)

type ErrNotAdded struct {
	Path string
}

func {{ .Function }}() (*fs.FS, error) {
{{ if eq .FS "DefaultFS" }}
	wfs := fs.DefaultFS
{{ else }}
	wfs := fs.NewFS()
{{ end }}
	var (
		size int64
		mode os.FileMode
		t    time.Time
	)
{{ range $index, $file := .Files }}
	size = int64({{ $file.Size }})
	mode = os.FileMode({{ $file.Mode }})
	t = time.Unix({{ $file.ModTime }}, 0)
	if !wfs.Add(fs.NewFileListing("{{ $file.Path }}", size, mode, t, []byte({{ $file.Data }}))) {
		return nil, &ErrNotAdded{Path: "$path"}
	}
{{ end }}
	return wfs, nil
}

func (e ErrNotAdded) Error() string {
	return fmt.Sprintf("Error adding file '%s'\n", e.Path)
}
`
<commit_msg>add support for adding build tags<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"text/template"
)

var (
	output    string
	pkg       string
	function  string
	fs        string
	buildTags string
	doFormat  bool

	tmpl = template.Must(template.New("go-template").Parse(goTemplate))
)

type templateData struct {
	Pkg      string
	Function string
	FS       string
	Tags     string

	Files []File
}

type File struct {
	Path    string
	Data    string
	Size    int64
	Mode    uint32
	ModTime int64
}

func main() {
	flag.Parse()
	if flag.NArg() == 0 {
		fmt.Fprintln(os.Stderr, "No input files/directories given")
		os.Exit(0)
	}

	if pkg == "" {
		fmt.Fprintln(os.Stderr, "No output file package given")
		os.Exit(0)
	}

	var out *os.File

	if output == "-" {
		out = os.Stdout
	} else {
		var err error
		if out, err = os.OpenFile(output, os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644); err != nil {
			fmt.Fprintf(os.Stderr, "Error opening '%s' for writing: %v\n", output, err)
			os.Exit(0)
		}

		defer out.Close()
	}

	files := []File{}

	for _, arg := range flag.Args() {
		recursive := false
		if strings.HasSuffix(arg, "/...") {
			recursive = true
			arg = arg[:len(arg)-4]
		}
		func() {
			file, err := os.Open(arg)
			defer file.Close()

			if err != nil {
				fmt.Fprintf(os.Stderr, "Error opening '%s': %v\n", arg, err)
				return
			}

			stat, err := file.Stat()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error stat-ing '%s': %v\n", arg, err)
				return
			}

			if stat.IsDir() {
				filepath.Walk(arg, func(path string, info os.FileInfo, err error) error {
					if info.IsDir() {
						if !recursive {
							return filepath.SkipDir
						}
					} else {
						if b, err := ioutil.ReadFile(path); err == nil {
							files = append(files, File{
								Path:    path,
								Data:    fmt.Sprintf("%q", b),
								Size:    info.Size(),
								ModTime: info.ModTime().Unix(),
								Mode:    uint32(info.Mode()),
							})
						} else {
							fmt.Fprintf(os.Stderr, "Error reading '%s': %v\n", path, err)
						}
					}

					return nil
				})
			} else {
				if b, err := ioutil.ReadAll(file); err == nil {
					files = append(files, File{
						Path:    arg,
						Data:    fmt.Sprintf("%q", b),
						Size:    stat.Size(),
						ModTime: stat.ModTime().Unix(),
						Mode:    uint32(stat.Mode()),
					})
				} else {
					fmt.Fprintf(os.Stderr, "Error reading '%s': %v\n", arg, err)
					return
				}

			}
		}()
	}

	buf := new(bytes.Buffer)
	err := tmpl.Execute(buf, templateData{Pkg: pkg, Function: function, FS: fs, Tags: buildTags, Files: files})
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error executing template: %v\n", err)
		os.Exit(0)
	}

	if doFormat {
		if b, err := format.Source(buf.Bytes()); err == nil {
			buf.Reset()
			if _, err = buf.Write(b); err != nil {
				fmt.Fprintf(os.Stderr, "Error writing formatted code to buffer: %v\n", err)
				os.Exit(0)
			}
		} else {
			fmt.Fprintf(os.Stderr, "Error formatting generated code: %v\n", err)
			os.Exit(0)
		}
	}

	if _, err := buf.WriteTo(out); err != nil {
		fmt.Fprintf(os.Stderr, "Error writing template result to output: %v\n", err)
		os.Exit(0)
	}
}

func init() {
	flag.StringVar(&output, "output", "-", "the output file")
	flag.StringVar(&pkg, "package", "main", "the package of the output go file")
	flag.StringVar(&function, "function", "addFiles", "the function that will add the files to the fs")
	flag.StringVar(&fs, "fs", "DefaultFS", "the fs object to be used. Will be created if different from the default")
	flag.StringVar(&buildTags, "build-tags", "", "optional build tags for the output code")
	flag.BoolVar(&doFormat, "format", false, "run the output go code through go/format")
}

const goTemplate = `
{{ if .Tags }}// +build {{ .Tags }}
{{ end }}
package {{ .Pkg }}

// DO NOT EDIT ** This file was generated with the webfw-fs tool ** DO NOT EDIT //

import (
	"fmt"
	"os"
	"time"

	"github.com/urandom/webfw/fs"
)

type ErrNotAdded struct {
	Path string
}

func {{ .Function }}() (*fs.FS, error) {
{{ if eq .FS "DefaultFS" }}
	wfs := fs.DefaultFS
{{ else }}
	wfs := fs.NewFS()
{{ end }}
	var (
		size int64
		mode os.FileMode
		t    time.Time
	)
{{ range $index, $file := .Files }}
	size = int64({{ $file.Size }})
	mode = os.FileMode({{ $file.Mode }})
	t = time.Unix({{ $file.ModTime }}, 0)
	if !wfs.Add(fs.NewFileListing("{{ $file.Path }}", size, mode, t, []byte({{ $file.Data }}))) {
		return nil, &ErrNotAdded{Path: "$path"}
	}
{{ end }}
	return wfs, nil
}

func (e ErrNotAdded) Error() string {
	return fmt.Sprintf("Error adding file '%s'\n", e.Path)
}
`
<|endoftext|>


<commit_before>// Copyright (c) 2016 Niklas Wolber
// This file is licensed under the MIT license.
// See the LICENSE file for more information.

package main

import (
	"flag"
	"fmt"
	"os"
	"time"
)

func config() (jobDir string, sshTTL time.Duration, file, logFile string, once bool) {
	const (
		jobDirDefault  = "."
		sshTTLDefault  = time.Minute * 10
		logFileDefault = ""
		fileDefault    = ""
		onceDefault    = false
	)

	flag.StringVar(&jobDir, "jobs", jobDirDefault, "Directory to watch for .job files")
	flag.DurationVar(&sshTTL, "sshTTL", sshTTLDefault, "Time until an unused SSH connection is closed")
	flag.StringVar(&file, "f", fileDefault, "Job file to execute. Disables automatic pick-up of job files.")
	flag.BoolVar(&once, "once", onceDefault, "Run job only once, regardless of the schedule. Only in combination with -f.")
	flag.StringVar(&logFile, "log", logFileDefault, "Log file")

	help := flag.Bool("help", false, "Display this help")
	config := flag.Bool("config", false, "Display current configuration")

	flag.Parse()

	if *help {
		flag.PrintDefaults()
		os.Exit(0)
	}

	if *config {
		fmt.Println("jobs  :", jobDir)
		fmt.Println("sshTTL:", sshTTL)
		os.Exit(0)
	}

	return
}
<commit_msg>Make -file the switch for one file mode<commit_after>// Copyright (c) 2016 Niklas Wolber
// This file is licensed under the MIT license.
// See the LICENSE file for more information.

package main

import (
	"flag"
	"fmt"
	"os"
	"time"
)

func config() (jobDir string, sshTTL time.Duration, file, logFile string, once bool) {
	const (
		jobDirDefault  = "."
		sshTTLDefault  = time.Minute * 10
		logFileDefault = ""
		fileDefault    = ""
		onceDefault    = false
	)

	flag.StringVar(&jobDir, "jobs", jobDirDefault, "Directory to watch for .job files")
	flag.DurationVar(&sshTTL, "sshTTL", sshTTLDefault, "Time until an unused SSH connection is closed")
	flag.StringVar(&file, "file", fileDefault, "Job file to execute. Disables automatic pick-up of job files.")
	flag.BoolVar(&once, "once", onceDefault, "Run job only once, regardless of the schedule. Only in combination with -f.")
	flag.StringVar(&logFile, "log", logFileDefault, "Log file")

	help := flag.Bool("help", false, "Display this help")
	config := flag.Bool("config", false, "Display current configuration")

	flag.Parse()

	if *help {
		flag.PrintDefaults()
		os.Exit(0)
	}

	if *config {
		fmt.Println("jobs  :", jobDir)
		fmt.Println("sshTTL:", sshTTL)
		os.Exit(0)
	}

	return
}
<|endoftext|>


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

// msr reads and writes msrs using a Forth interpreter on argv
//
// Synopsis:
//     To see what is available:
//     msr words
//
// Description:
//     msr provides a set of Forth words that let you manage MSRs.
//     You can add new ones of your own.
//     For a start, it provides some pre-defined words for well-known MSRs
//
//     push a [] of MSR names and the 0x3a register on the stack
//     IA32_FEATURE_CONTROL -- equivalent to * msr 0x3a reg
//     The next two commands use IA32_FEATURE_CONTROL:
//     READ_IA32_FEATURE_CONTROL -- equivalent to IA32_FEATURE_CONTROL rd
//     LOCK IA32_FEATURE_CONTROL -- equivalent to IA32_FEATURE_CONTROL rd IA32_FEATURE_CONTROL 1 u64 or wr
//     e.g.
//./msr IA32_FEATURE_CONTROL
// [[/dev/cpu/0/msr /dev/cpu/1/msr /dev/cpu/2/msr /dev/cpu/3/msr] 58]
//
//     As a special convenience, we have two useful cases:
//     r glob register -- read the MSR 'register' from cores matching 'glob'
//     w glob register value -- write the value to 'register' on all cores matching 'glob'
//
// Examples:
//     Show the IA32 feature MSR on all cores
//     sudo fio READ_IA32_FEATURE_CONTROL
//     [[5 5 5 5]]
//     lock the registers
//     sudo fio LOCK_IA32_FEATURE_CONTROL
//     Just see it one core 0 and 1
//     sudo ./fio '[01]' msr 0x3a reg rd
//     [[5 5]]
package main

import (
	"flag"
	"fmt"
	"log"
	"strconv"

	"github.com/u-root/u-root/pkg/forth"
	"github.com/u-root/u-root/pkg/msr"
)

// let's just do MSRs for now

var (
	debug = flag.Bool("d", false, "debug messages")
	words = []struct {
		name string
		w    []forth.Cell
	}{
		// Architectural MSR. All systems.
		// Enables features like VMX.
		{name: "MSR_IA32_FEATURE_CONTROL", w: []forth.Cell{"'*", "cpu", "0x3a", "reg"}},
		{name: "READ_MSR_IA32_FEATURE_CONTROL", w: []forth.Cell{"MSR_IA32_FEATURE_CONTROL", "rd"}},
		{name: "LOCK_MSR_IA32_FEATURE_CONTROL", w: []forth.Cell{"MSR_IA32_FEATURE_CONTROL", "READ_MSR_IA32_FEATURE_CONTROL", "1", "u64", "or", "wr"}},
		// Silvermont, Airmont, Nehalem...
		// Controls Processor C States.
		{name: "MSR_PKG_CST_CONFIG_CONTROL", w: []forth.Cell{"'*", "cpu", "0xe2", "reg"}},
		{name: "READ_MSR_PKG_CST_CONFIG_CONTROL", w: []forth.Cell{"MSR_PKG_CST_CONFIG_CONTROL", "rd"}},
		{name: "LOCK_MSR_PKG_CST_CONFIG_CONTROL", w: []forth.Cell{"MSR_PKG_CST_CONFIG_CONTROL", "READ_MSR_PKG_CST_CONFIG_CONTROL", uint64(1 << 15), "or", "wr"}},
		// Westmere onwards.
		// Note that this turns on AES instructions, however
		// 3 will turn off AES until reset.
		{name: "MSR_FEATURE_CONFIG", w: []forth.Cell{"'*", "cpu", "0x13c", "reg"}},
		{name: "READ_MSR_FEATURE_CONFIG", w: []forth.Cell{"MSR_FEATURE_CONFIG", "rd"}},
		{name: "LOCK_MSR_FEATURE_CONFIG", w: []forth.Cell{"MSR_FEATURE_CONFIG", "READ_MSR_FEATURE_CONFIG", uint64(1 << 0), "or", "wr"}},
		// Goldmont, SandyBridge
		// Controls DRAM power limits. See Intel SDM
		{name: "MSR_DRAM_POWER_LIMIT", w: []forth.Cell{"'*", "cpu", "0x618", "reg"}},
		{name: "READ_MSR_DRAM_POWER_LIMIT", w: []forth.Cell{"MSR_DRAM_POWER_LIMIT", "rd"}},
		{name: "LOCK_MSR_DRAM_POWER_LIMIT", w: []forth.Cell{"MSR_DRAM_POWER_LIMIT", "READ_MSR_DRAM_POWER_LIMIT", uint64(1 << 31), "or", "wr"}},
		// IvyBridge Onwards.
		// Not much information in the SDM, seems to control power limits
		{name: "MSR_CONFIG_TDP_CONTROL", w: []forth.Cell{"'*", "cpu", "0xe2", "reg"}},
		{name: "READ_MSR_CONFIG_TDP_CONTROL", w: []forth.Cell{"MSR_CONFIG_TDP_CONTROL", "rd"}},
		{name: "LOCK_MSR_CONFIG_TDP_CONTROL", w: []forth.Cell{"MSR_CONFIG_TDP_CONTROL", "READ_MSR_CONFIG_TDP_CONTROL", uint64(1 << 31), "or", "wr"}},
		// Architectural MSR. All systems.
		// This is the actual spelling of the MSR in the manual.
		// Controls availability of silicon debug interfaces
		{name: "IA32_DEBUG_INTERFACE", w: []forth.Cell{"'*", "cpu", "0xe2", "reg"}},
		{name: "READ_IA32_DEBUG_INTERFACE", w: []forth.Cell{"IA32_DEBUG_INTERFACE", "rd"}},
		{name: "LOCK_IA32_DEBUG_INTERFACE", w: []forth.Cell{"IA32_DEBUG_INTERFACE", "READ_IA32_DEBUG_INTERFACE", uint64(1 << 15), "or", "wr"}},
		// Locks all known msrs to lock
		{name: "LOCK_KNOWN_MSRS", w: []forth.Cell{"LOCK_MSR_IA32_FEATURE_CONTROL", "LOCK_MSR_PKG_CST_CONFIG_CONTROL", "LOCK_MSR_FEATURE_CONFIG", "LOCK_MSR_DRAM_POWER_LIMIT", "LOCK_MSR_CONFIG_TDP_CONTROL", "LOCK_IA32_DEBUG_INTERFACE"}},
	}
	ops = []struct {
		name string
		op   forth.Op
	}{
		{name: "cpu", op: cpus},
		{name: "reg", op: reg},
		{name: "u64", op: u64},
		{name: "rd", op: rd},
		{name: "wr", op: wr},
		{name: "swr", op: swr},
		{name: "and", op: and},
		{name: "or", op: or},
	}
)

// The great panic discussion.
// Rob has given talks on using panic for parsers.
// I have talked to Russ about using panic for parsers.
// Short form: it's ok. In general, don't panic.
// But parsers are special: using panic
// in a parser makes the code tons cleaner.

// Note that if any type asserts fail the forth interpret loop catches
// it. It also catches stack underflow, all that stuff.
func cpus(f forth.Forth) {
	forth.Debug("cpu")
	g := f.Pop().(string)
	c, errs := msr.GlobCPUs(g)
	if errs != nil {
		panic(fmt.Sprintf("%v", errs))
	}
	forth.Debug("CPUs are %v", c)
	f.Push(c)
}

func reg(f forth.Forth) {
	n, err := strconv.ParseUint(f.Pop().(string), 0, 32)
	if err != nil {
		panic(fmt.Sprintf("%v", err))
	}
	f.Push(msr.MSR(n))
}

func u64(f forth.Forth) {
	n, err := strconv.ParseUint(f.Pop().(string), 0, 64)
	if err != nil {
		panic(fmt.Sprintf("%v", err))
	}
	f.Push(uint64(n))
}

func rd(f forth.Forth) {
	r := f.Pop().(msr.MSR)
	c := f.Pop().(msr.CPUs)
	forth.Debug("rd: cpus %v, msr %v", c, r)
	data, errs := r.Read(c)
	forth.Debug("data %v errs %v", data, errs)
	if errs != nil {
		panic(fmt.Sprintf("%v", errs))
	}
	f.Push(data)
}

// the standard msr api takes one value for all msrs.
// That's arguably substandard. We're going to require
// a []uint64. It's naive to expect every core to have
// exactly the same msr values for each msr in our
// modern world.
// If you're determined to write a fixed value, the same
// for all, it's easy:
// fio "'"* msr 0x3a reg rd 0 val and your-value new-val val or wr
// Then you'll have a fixed value.
func wr(f forth.Forth) {
	v := f.Pop().([]uint64)
	r := f.Pop().(msr.MSR)
	c := f.Pop().(msr.CPUs)
	forth.Debug("wr: cpus %v, msr %v, values %v", c, r, v)
	errs := r.Write(c, v...)
	forth.Debug("errs %v", errs)
	if errs != nil {
		f.Push(errs)
	}
}

// We had been counting on doing a rd, which would produce a nice
// []u64 at TOS which we could use in a write. Problem is, some MSRs
// can not be read. There are write-only MSRs.  This really
// complicates the picture: we can't just read them, change them, and
// write them; we would not even know if reading is side-effect free.
//
// Write now accepts a single value
func swr(f forth.Forth) {
	v := f.Pop().(uint64)
	r := f.Pop().(msr.MSR)
	c := f.Pop().(msr.CPUs)

	forth.Debug("swr: cpus %v, msr %v, %v", c, r, v)
	errs := r.Write(c, v)
	forth.Debug("errs %v", errs)
	if errs != nil {
		f.Push(errs)
	}
}

func and(f forth.Forth) {
	v := f.Pop().(uint64)
	m := f.Pop().([]uint64)
	forth.Debug("and: %v(%T) %v(%T)", m, m, v, v)
	for i := range m {
		m[i] &= v
	}
	f.Push(m)
}

func or(f forth.Forth) {
	v := f.Pop().(uint64)
	m := f.Pop().([]uint64)
	forth.Debug("or: %v(%T) %v(%T)", m, m, v, v)
	for i := range m {
		m[i] |= v
	}
	f.Push(m)
}

func main() {
	flag.Parse()
	if *debug {
		forth.Debug = log.Printf
	}

	// TODO: probably can do this by just having two passes, and write
	// in the first pass is a no op. Which will fail to catch the problem
	// of read-only and write-only MSRs but there's only so much you can do.
	//
	// To avoid the command list from being partially executed when the
	// args fail to parse, queue them up and run all at once at the end.
	//queue := []func(){} etc. etc.
	f := forth.New()
	for _, o := range ops {
		forth.Putop(o.name, o.op)
	}
	for _, w := range words {
		forth.NewWord(f, w.name, w.w[0], w.w[1:]...)
	}
	a := flag.Args()
	// If the first arg is r or w, we're going to assume they're not doing Forth.
	// It is too confusing otherwise if they type a wrong r or w command and
	// see the Forth stack and nothing else.
	switch a[0] {
	case "r":
		if len(a) != 3 {
			log.Fatal("Usage for r: r <msr-glob> <register>")
		}
		// Because the msr arg is a glob and may have things like * in it (* being the
		// most common) gratuitiously add a Forth ' before it (i.e. quote it).
		if err := forth.EvalString(f, fmt.Sprintf("'%s msr %s reg rd", a[1], a[2])); err != nil {
			log.Fatal(err)
		}
	case "w":
		if len(a) != 4 {
			log.Fatal("Usage for w: w <msr-glob> <register> <value>")
		}
		// Because the msr arg is a glob and may have things like * in it (* being the
		// most common) gratuitiously add a Forth ' before it (i.e. quote it).
		if err := forth.EvalString(f, fmt.Sprintf("'%s msr %s reg %s u64 swr", a[1], a[2], a[3])); err != nil {
			log.Fatal(err)
		}
	case "lock":
		if len(a) != 4 {
			log.Fatal("Usage for lock: lock <msr-glob> <register> <bit>")
		}
		if err := forth.EvalString(f, fmt.Sprintf("'%s msr %s reg '%s msr %s reg rd %s u64 or wr", a[1], a[2], a[1], a[2], a[3])); err != nil {
			log.Fatal(err)
		}
	default:
		for _, a := range flag.Args() {
			if err := forth.EvalString(f, a); err != nil {
				log.Fatal(err)
			}
			forth.Debug("%vOK\n", f.Stack())
		}
	}
	// special case: if the length of stack is 1, just print out stack[0].
	s := f.Stack()
	if len(s) == 1 {
		fmt.Printf("%v\n", s[0])
	} else {
		fmt.Printf("%v\n", s)
	}
}
<commit_msg>Fix dumb MSR bug<commit_after>// Copyright 2010-2019 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// msr reads and writes msrs using a Forth interpreter on argv
//
// Synopsis:
//     To see what is available:
//     msr words
//
// Description:
//     msr provides a set of Forth words that let you manage MSRs.
//     You can add new ones of your own.
//     For a start, it provides some pre-defined words for well-known MSRs
//
//     push a [] of MSR names and the 0x3a register on the stack
//     IA32_FEATURE_CONTROL -- equivalent to * msr 0x3a reg
//     The next two commands use IA32_FEATURE_CONTROL:
//     READ_IA32_FEATURE_CONTROL -- equivalent to IA32_FEATURE_CONTROL rd
//     LOCK IA32_FEATURE_CONTROL -- equivalent to IA32_FEATURE_CONTROL rd IA32_FEATURE_CONTROL 1 u64 or wr
//     e.g.
//./msr IA32_FEATURE_CONTROL
// [[/dev/cpu/0/msr /dev/cpu/1/msr /dev/cpu/2/msr /dev/cpu/3/msr] 58]
//
//     As a special convenience, we have two useful cases:
//     r glob register -- read the MSR 'register' from cores matching 'glob'
//     w glob register value -- write the value to 'register' on all cores matching 'glob'
//
// Examples:
//     Show the IA32 feature MSR on all cores
//     sudo fio READ_IA32_FEATURE_CONTROL
//     [[5 5 5 5]]
//     lock the registers
//     sudo fio LOCK_IA32_FEATURE_CONTROL
//     Just see it one core 0 and 1
//     sudo ./fio '[01]' msr 0x3a reg rd
//     [[5 5]]
package main

import (
	"flag"
	"fmt"
	"log"
	"strconv"

	"github.com/u-root/u-root/pkg/forth"
	"github.com/u-root/u-root/pkg/msr"
)

// let's just do MSRs for now

var (
	debug = flag.Bool("d", false, "debug messages")
	words = []struct {
		name string
		w    []forth.Cell
	}{
		// Architectural MSR. All systems.
		// Enables features like VMX.
		{name: "MSR_IA32_FEATURE_CONTROL", w: []forth.Cell{"'*", "cpu", "0x3a", "reg"}},
		{name: "READ_MSR_IA32_FEATURE_CONTROL", w: []forth.Cell{"MSR_IA32_FEATURE_CONTROL", "rd"}},
		{name: "LOCK_MSR_IA32_FEATURE_CONTROL", w: []forth.Cell{"MSR_IA32_FEATURE_CONTROL", "READ_MSR_IA32_FEATURE_CONTROL", "1", "u64", "or", "wr"}},
		// Silvermont, Airmont, Nehalem...
		// Controls Processor C States.
		{name: "MSR_PKG_CST_CONFIG_CONTROL", w: []forth.Cell{"'*", "cpu", "0xe2", "reg"}},
		{name: "READ_MSR_PKG_CST_CONFIG_CONTROL", w: []forth.Cell{"MSR_PKG_CST_CONFIG_CONTROL", "rd"}},
		{name: "LOCK_MSR_PKG_CST_CONFIG_CONTROL", w: []forth.Cell{"MSR_PKG_CST_CONFIG_CONTROL", "READ_MSR_PKG_CST_CONFIG_CONTROL", uint64(1 << 15), "or", "wr"}},
		// Westmere onwards.
		// Note that this turns on AES instructions, however
		// 3 will turn off AES until reset.
		{name: "MSR_FEATURE_CONFIG", w: []forth.Cell{"'*", "cpu", "0x13c", "reg"}},
		{name: "READ_MSR_FEATURE_CONFIG", w: []forth.Cell{"MSR_FEATURE_CONFIG", "rd"}},
		{name: "LOCK_MSR_FEATURE_CONFIG", w: []forth.Cell{"MSR_FEATURE_CONFIG", "READ_MSR_FEATURE_CONFIG", uint64(1 << 0), "or", "wr"}},
		// Goldmont, SandyBridge
		// Controls DRAM power limits. See Intel SDM
		{name: "MSR_DRAM_POWER_LIMIT", w: []forth.Cell{"'*", "cpu", "0x618", "reg"}},
		{name: "READ_MSR_DRAM_POWER_LIMIT", w: []forth.Cell{"MSR_DRAM_POWER_LIMIT", "rd"}},
		{name: "LOCK_MSR_DRAM_POWER_LIMIT", w: []forth.Cell{"MSR_DRAM_POWER_LIMIT", "READ_MSR_DRAM_POWER_LIMIT", uint64(1 << 31), "or", "wr"}},
		// IvyBridge Onwards.
		// Not much information in the SDM, seems to control power limits
		{name: "MSR_CONFIG_TDP_CONTROL", w: []forth.Cell{"'*", "cpu", "0xe2", "reg"}},
		{name: "READ_MSR_CONFIG_TDP_CONTROL", w: []forth.Cell{"MSR_CONFIG_TDP_CONTROL", "rd"}},
		{name: "LOCK_MSR_CONFIG_TDP_CONTROL", w: []forth.Cell{"MSR_CONFIG_TDP_CONTROL", "READ_MSR_CONFIG_TDP_CONTROL", uint64(1 << 31), "or", "wr"}},
		// Architectural MSR. All systems.
		// This is the actual spelling of the MSR in the manual.
		// Controls availability of silicon debug interfaces
		{name: "IA32_DEBUG_INTERFACE", w: []forth.Cell{"'*", "cpu", "0xe2", "reg"}},
		{name: "READ_IA32_DEBUG_INTERFACE", w: []forth.Cell{"IA32_DEBUG_INTERFACE", "rd"}},
		{name: "LOCK_IA32_DEBUG_INTERFACE", w: []forth.Cell{"IA32_DEBUG_INTERFACE", "READ_IA32_DEBUG_INTERFACE", uint64(1 << 15), "or", "wr"}},
		// Locks all known msrs to lock
		{name: "LOCK_KNOWN_MSRS", w: []forth.Cell{"LOCK_MSR_IA32_FEATURE_CONTROL", "LOCK_MSR_PKG_CST_CONFIG_CONTROL", "LOCK_MSR_FEATURE_CONFIG", "LOCK_MSR_DRAM_POWER_LIMIT", "LOCK_MSR_CONFIG_TDP_CONTROL", "LOCK_IA32_DEBUG_INTERFACE"}},
	}
	ops = []struct {
		name string
		op   forth.Op
	}{
		{name: "cpu", op: cpus},
		{name: "reg", op: reg},
		{name: "u64", op: u64},
		{name: "rd", op: rd},
		{name: "wr", op: wr},
		{name: "swr", op: swr},
		{name: "and", op: and},
		{name: "or", op: or},
	}
)

// The great panic discussion.
// Rob has given talks on using panic for parsers.
// I have talked to Russ about using panic for parsers.
// Short form: it's ok. In general, don't panic.
// But parsers are special: using panic
// in a parser makes the code tons cleaner.

// Note that if any type asserts fail the forth interpret loop catches
// it. It also catches stack underflow, all that stuff.
func cpus(f forth.Forth) {
	forth.Debug("cpu")
	g := f.Pop().(string)
	c, errs := msr.GlobCPUs(g)
	if errs != nil {
		panic(fmt.Sprintf("%v", errs))
	}
	forth.Debug("CPUs are %v", c)
	f.Push(c)
}

func reg(f forth.Forth) {
	n, err := strconv.ParseUint(f.Pop().(string), 0, 32)
	if err != nil {
		panic(fmt.Sprintf("%v", err))
	}
	f.Push(msr.MSR(n))
}

func u64(f forth.Forth) {
	n, err := strconv.ParseUint(f.Pop().(string), 0, 64)
	if err != nil {
		panic(fmt.Sprintf("%v", err))
	}
	f.Push(uint64(n))
}

func rd(f forth.Forth) {
	r := f.Pop().(msr.MSR)
	c := f.Pop().(msr.CPUs)
	forth.Debug("rd: cpus %v, msr %v", c, r)
	data, errs := r.Read(c)
	forth.Debug("data %v errs %v", data, errs)
	if errs != nil {
		panic(fmt.Sprintf("%v", errs))
	}
	f.Push(data)
}

// the standard msr api takes one value for all msrs.
// That's arguably substandard. We're going to require
// a []uint64. It's naive to expect every core to have
// exactly the same msr values for each msr in our
// modern world.
// If you're determined to write a fixed value, the same
// for all, it's easy:
// fio "'"* msr 0x3a reg rd 0 val and your-value new-val val or wr
// Then you'll have a fixed value.
func wr(f forth.Forth) {
	v := f.Pop().([]uint64)
	r := f.Pop().(msr.MSR)
	c := f.Pop().(msr.CPUs)
	forth.Debug("wr: cpus %v, msr %v, values %v", c, r, v)
	errs := r.Write(c, v...)
	forth.Debug("errs %v", errs)
	if errs != nil {
		f.Push(errs)
	}
}

// We had been counting on doing a rd, which would produce a nice
// []u64 at TOS which we could use in a write. Problem is, some MSRs
// can not be read. There are write-only MSRs.  This really
// complicates the picture: we can't just read them, change them, and
// write them; we would not even know if reading is side-effect free.
//
// Write now accepts a single value
func swr(f forth.Forth) {
	v := f.Pop().(uint64)
	r := f.Pop().(msr.MSR)
	c := f.Pop().(msr.CPUs)

	forth.Debug("swr: cpus %v, msr %v, %v", c, r, v)
	errs := r.Write(c, v)
	forth.Debug("errs %v", errs)
	if errs != nil {
		f.Push(errs)
	}
}

func and(f forth.Forth) {
	v := f.Pop().(uint64)
	m := f.Pop().([]uint64)
	forth.Debug("and: %v(%T) %v(%T)", m, m, v, v)
	for i := range m {
		m[i] &= v
	}
	f.Push(m)
}

func or(f forth.Forth) {
	v := f.Pop().(uint64)
	m := f.Pop().([]uint64)
	forth.Debug("or: %v(%T) %v(%T)", m, m, v, v)
	for i := range m {
		m[i] |= v
	}
	f.Push(m)
}

func main() {
	flag.Parse()
	if *debug {
		forth.Debug = log.Printf
	}

	// TODO: probably can do this by just having two passes, and write
	// in the first pass is a no op. Which will fail to catch the problem
	// of read-only and write-only MSRs but there's only so much you can do.
	//
	// To avoid the command list from being partially executed when the
	// args fail to parse, queue them up and run all at once at the end.
	//queue := []func(){} etc. etc.
	f := forth.New()
	for _, o := range ops {
		forth.Putop(o.name, o.op)
	}
	for _, w := range words {
		forth.NewWord(f, w.name, w.w[0], w.w[1:]...)
	}
	a := flag.Args()
	// If the first arg is r or w, we're going to assume they're not doing Forth.
	// It is too confusing otherwise if they type a wrong r or w command and
	// see the Forth stack and nothing else.
	switch a[0] {
	case "r":
		if len(a) != 3 {
			log.Fatal("Usage for r: r <msr-glob> <register>")
		}
		// Because the msr arg is a glob and may have things like * in it (* being the
		// most common) gratuitiously add a Forth ' before it (i.e. quote it).
		if err := forth.EvalString(f, fmt.Sprintf("'%s cpu %s reg rd", a[1], a[2])); err != nil {
			log.Fatal(err)
		}
	case "w":
		if len(a) != 4 {
			log.Fatal("Usage for w: w <msr-glob> <register> <value>")
		}
		// Because the msr arg is a glob and may have things like * in it (* being the
		// most common) gratuitiously add a Forth ' before it (i.e. quote it).
		if err := forth.EvalString(f, fmt.Sprintf("'%s cpu %s reg %s u64 swr", a[1], a[2], a[3])); err != nil {
			log.Fatal(err)
		}
	case "lock":
		if len(a) != 4 {
			log.Fatal("Usage for lock: lock <msr-glob> <register> <bit>")
		}
		if err := forth.EvalString(f, fmt.Sprintf("'%s cpu %s reg '%s msr %s reg rd %s u64 or wr", a[1], a[2], a[1], a[2], a[3])); err != nil {
			log.Fatal(err)
		}
	default:
		for _, a := range flag.Args() {
			if err := forth.EvalString(f, a); err != nil {
				log.Fatal(err)
			}
			forth.Debug("%vOK\n", f.Stack())
		}
	}
	// special case: if the length of stack is 1, just print out stack[0].
	s := f.Stack()
	if len(s) == 1 {
		fmt.Printf("%v\n", s[0])
	} else {
		fmt.Printf("%v\n", s)
	}
}
<|endoftext|>


<commit_before>package val

import (
	"bytes"
	"context"
	"database/sql"
	"github.com/ghts/ghts/lib"
	"strconv"
	"time"
)

func F내재가치_정보_모음_DB읽기(db *sql.DB) (s *S내재가치_정보_모음, 에러 error) {
	s = New내재가치_정보_모음()
	에러 = s.DB읽기(db)

	return s, 에러
}

func New내재가치_정보_모음() *S내재가치_정보_모음 {
	s := new(S내재가치_정보_모음)
	s.M저장소 = make(map[string]*S내재가치_정보)

	return s
}

type S내재가치_정보_모음 struct {
	M저장소 map[string]*S내재가치_정보
}

func (s *S내재가치_정보_모음) G종목별_최신_정보(종목코드 string) *S내재가치_정보 {
	for 연도 := lib.F지금().Year(); 연도 > 2015; 연도-- {
		for 월 := 12; 월 > 0; 월-- {
			키 := 종목코드 + lib.F2문자열(연도*100+월)
			if 값, 존재함 := s.M저장소[키]; 존재함 {
				return 값
			}
		}
	}

	return nil
}

func (s *S내재가치_정보_모음) G종목별_차최신_정보(종목코드 string) *S내재가치_정보 {
	차최신_연도 := s.G종목별_최신_정보(종목코드).S내재가치_식별정보.G일자2().Year() - 1
	for 연도 := 차최신_연도; 연도 > 2015; 연도-- {
		for 월 := 12; 월 > 0; 월-- {
			키 := 종목코드 + lib.F2문자열(연도*100+월)
			if 값, 존재함 := s.M저장소[키]; 존재함 {
				return 값
			}
		}
	}

	return nil
}

func (s *S내재가치_정보_모음) G종목별_차차최신_정보(종목코드 string) *S내재가치_정보 {
	차차최신_연도 := s.G종목별_최신_정보(종목코드).S내재가치_식별정보.G일자2().Year() - 2
	for 연도 := 차차최신_연도; 연도 > 2015; 연도-- {
		for 월 := 12; 월 > 0; 월-- {
			키 := 종목코드 + lib.F2문자열(연도*100+월)
			if 값, 존재함 := s.M저장소[키]; 존재함 {
				return 값
			}
		}
	}

	return nil
}

func (s *S내재가치_정보_모음) S상장주식수_업데이트(db *sql.DB, 종목코드 string, 수량 int64) error {
	if 값 := s.G종목별_최신_정보(종목코드); 값 != nil {
		값.M상장주식수 = float64(수량)
		return F내재가치_정보_모음_DB저장(db, []*S내재가치_정보{값})
	}

	return nil
}

func (s *S내재가치_정보_모음) S파일_읽기(파일명 string) error {
	return lib.JSON_파일_읽기(파일명, s)
}

func (s *S내재가치_정보_모음) S파일_저장(파일명 string) error {
	return lib.JSON_파일_저장(s, 파일명)
}

func (s *S내재가치_정보_모음) DB읽기(db *sql.DB) (에러 error) {
	defer lib.S예외처리{M에러: &에러, M함수: func() { s = nil }}.S실행()

	lib.F확인(F내재가치_정보_테이블_생성(db))

	SQL := new(bytes.Buffer)
	SQL.WriteString("SELECT code, date, json ")
	SQL.WriteString("FROM fundamental_data ")
	SQL.WriteString("ORDER BY code, date")

	stmt, 에러 := db.Prepare(SQL.String())
	lib.F확인(에러)
	defer stmt.Close()

	rows, 에러 := stmt.Query()
	lib.F확인(에러)
	defer rows.Close()

	for rows.Next() {
		var code string
		var date time.Time
		var json string
		var 값 *S내재가치_정보

		lib.F확인(rows.Scan(&code, &date, &json))

		if 에러 = lib.F디코딩(lib.JSON, []byte(json), &값); 에러 != nil {
			lib.F문자열_출력("%v %v : 디코딩 에러\n%v", code, date.Format(lib.P일자_형식), 에러)
			continue
		} else if 값 == nil {
			lib.F문자열_출력("%v %v : nil 값", code, date.Format(lib.P일자_형식))
			continue
		} else if 값.S내재가치_식별정보 == nil {
			lib.F문자열_출력("%v %v : nil 식별정보", code, date.Format(lib.P일자_형식))
			continue
		}

		s.M저장소[값.G키()] = 값
	}

	return
}

func (s *S내재가치_정보_모음) DB저장(db *sql.DB) error {
	// 맵 데이터를 슬라이스 형태로 변환
	값_모음 := make([]*S내재가치_정보, len(s.M저장소))

	i := 0
	for _, 값 := range s.M저장소 {
		값_모음[i] = 값
		i++
	}

	// 슬라이스 데이터 DB에 저장
	return F내재가치_정보_모음_DB저장(db, 값_모음)
}

type S내재가치_정보 struct {
	*S내재가치_식별정보
	*S재무제표_정보_내용
	*S재무비율_정보_내용
	*S투자지표_정보_내용
}

type S재무제표_정보 struct {
	*S내재가치_식별정보
	*S재무제표_정보_내용
}

type S재무비율_정보 struct {
	*S내재가치_식별정보
	*S재무비율_정보_내용
}

type S투자지표_정보 struct {
	*S내재가치_식별정보
	*S투자지표_정보_내용
}

func New내재가치_식별정보(종목코드 string, 연도 uint16, 월 uint8) *S내재가치_식별정보 {
	s := new(S내재가치_식별정보)
	s.M종목코드 = 종목코드
	s.S연도_월(연도, 월)

	return s
}

type S내재가치_식별정보 struct {
	M종목코드 string
	M연도_월 uint32
}

func (s S내재가치_식별정보) G키() string {
	return s.M종목코드 + strconv.Itoa(int(s.M연도_월))
}

func (s *S내재가치_식별정보) S연도_월(연도 uint16, 월 uint8) {
	s.M연도_월 = uint32(연도)*100 + uint32(월)
}

func (s S내재가치_식별정보) G일자() uint32 {
	return s.M연도_월*100 + 1
}

func (s S내재가치_식별정보) G일자2() time.Time {
	return lib.F2포맷된_일자_단순형("20060102", strconv.Itoa(int(s.M연도_월*100+1)))
}

type S재무제표_정보_내용 struct {
	M매출액      float64
	M매출총이익    float64
	M영업이익     float64
	M세전계속사업이익 float64
	M당기순이익    float64
	M자산       float64
	M유동자산     float64 // NCAV 전략에 필요함.
	M현금및현금성자산 float64 // EV계산에 필요함.
	M부채       float64
	M유동부채     float64 // NCAV 전략에 필요함.
	M자본       float64
	M영업_현금흐름  float64
	M투자_현금흐름  float64
	M재무_현금흐름  float64
	M현금_증가    float64
}

type S재무비율_정보_내용 struct {
	M유동_비율     float64
	M당좌_비율     float64
	M부채_비율     float64
	M이자보상배율    float64
	M매출총이익율    float64
	M세전계속사업이익율 float64
	M영업이익율     float64
	EBITDA마진율  float64
	ROA        float64
	ROE        float64
	ROIC       float64
	M총자산회전율    float64
	M총부채회전율    float64
	M총자본회전율    float64
	M순운전자본회전율  float64
}

type S투자지표_정보_내용 struct {
	M상장주식수    float64
	EPS       float64
	EBITDAPS  float64
	CFPS      float64
	SPS       float64
	BPS       float64
	DPS       float64
	M배당성향     float64
	M총현금흐름    float64
	M세후영업이익   float64
	M유무형자산상각비 float64
	M총투자      float64
	FCFF      float64
}

func F내재가치_정보_모음_DB저장(db *sql.DB, 값_모음 []*S내재가치_정보) (에러 error) {
	var tx *sql.Tx
	defer lib.S예외처리{M에러: &에러, M함수: func() { lib.F조건부_실행(tx != nil, tx.Rollback) }}.S실행()

	lib.F확인(F내재가치_정보_테이블_생성(db))

	SQL생성 := new(bytes.Buffer)
	SQL생성.WriteString("INSERT IGNORE INTO fundamental_data (")
	SQL생성.WriteString("code,")
	SQL생성.WriteString("date,")
	SQL생성.WriteString("json")
	SQL생성.WriteString(") VALUES (?,?,'')")

	SQL수정 := new(bytes.Buffer)
	SQL수정.WriteString("UPDATE fundamental_data SET json=? ")
	SQL수정.WriteString("WHERE code=? AND date=?")

	txOpts := new(sql.TxOptions)
	txOpts.Isolation = sql.LevelDefault
	txOpts.ReadOnly = false

	tx, 에러 = db.BeginTx(context.TODO(), txOpts)
	lib.F확인(에러)

	stmt생성, 에러 := tx.Prepare(SQL생성.String())
	lib.F확인(에러)
	defer stmt생성.Close()

	stmt수정, 에러 := tx.Prepare(SQL수정.String())
	lib.F확인(에러)
	defer stmt수정.Close()

	for _, 값 := range 값_모음 {
		json, 에러 := lib.F인코딩(lib.JSON, 값)
		lib.F확인(에러)

		_, 에러 = stmt생성.Exec(값.M종목코드, 값.G일자())
		lib.F확인(에러)

		_, 에러 = stmt수정.Exec(string(json), 값.M종목코드, 값.G일자())
		lib.F확인(에러)
	}

	tx.Commit()

	return nil
}

func F내재가치_정보_테이블_생성(db *sql.DB) error {
	SQL := new(bytes.Buffer)
	SQL.WriteString("CREATE TABLE IF NOT EXISTS fundamental_data (")
	SQL.WriteString("code CHAR(8) NOT NULL,")
	SQL.WriteString("date DATE NOT NULL,")
	SQL.WriteString("json TEXT NOT NULL,")
	SQL.WriteString("CONSTRAINT PRIMARY KEY (code,date)")
	SQL.WriteString(")")

	_, 에러 := db.Exec(SQL.String())

	return 에러
}
<commit_msg>NCAV 전략에서 유동부채는 필요없다는 것을 뒤늦게 깨달음.<commit_after>package val

import (
	"bytes"
	"context"
	"database/sql"
	"github.com/ghts/ghts/lib"
	"strconv"
	"time"
)

func F내재가치_정보_모음_DB읽기(db *sql.DB) (s *S내재가치_정보_모음, 에러 error) {
	s = New내재가치_정보_모음()
	에러 = s.DB읽기(db)

	return s, 에러
}

func New내재가치_정보_모음() *S내재가치_정보_모음 {
	s := new(S내재가치_정보_모음)
	s.M저장소 = make(map[string]*S내재가치_정보)

	return s
}

type S내재가치_정보_모음 struct {
	M저장소 map[string]*S내재가치_정보
}

func (s *S내재가치_정보_모음) G종목별_최신_정보(종목코드 string) *S내재가치_정보 {
	for 연도 := lib.F지금().Year(); 연도 > 2015; 연도-- {
		for 월 := 12; 월 > 0; 월-- {
			키 := 종목코드 + lib.F2문자열(연도*100+월)
			if 값, 존재함 := s.M저장소[키]; 존재함 {
				return 값
			}
		}
	}

	return nil
}

func (s *S내재가치_정보_모음) G종목별_차최신_정보(종목코드 string) *S내재가치_정보 {
	차최신_연도 := s.G종목별_최신_정보(종목코드).S내재가치_식별정보.G일자2().Year() - 1
	for 연도 := 차최신_연도; 연도 > 2015; 연도-- {
		for 월 := 12; 월 > 0; 월-- {
			키 := 종목코드 + lib.F2문자열(연도*100+월)
			if 값, 존재함 := s.M저장소[키]; 존재함 {
				return 값
			}
		}
	}

	return nil
}

func (s *S내재가치_정보_모음) G종목별_차차최신_정보(종목코드 string) *S내재가치_정보 {
	차차최신_연도 := s.G종목별_최신_정보(종목코드).S내재가치_식별정보.G일자2().Year() - 2
	for 연도 := 차차최신_연도; 연도 > 2015; 연도-- {
		for 월 := 12; 월 > 0; 월-- {
			키 := 종목코드 + lib.F2문자열(연도*100+월)
			if 값, 존재함 := s.M저장소[키]; 존재함 {
				return 값
			}
		}
	}

	return nil
}

func (s *S내재가치_정보_모음) S상장주식수_업데이트(db *sql.DB, 종목코드 string, 수량 int64) error {
	if 값 := s.G종목별_최신_정보(종목코드); 값 != nil {
		값.M상장주식수 = float64(수량)
		return F내재가치_정보_모음_DB저장(db, []*S내재가치_정보{값})
	}

	return nil
}

func (s *S내재가치_정보_모음) S파일_읽기(파일명 string) error {
	return lib.JSON_파일_읽기(파일명, s)
}

func (s *S내재가치_정보_모음) S파일_저장(파일명 string) error {
	return lib.JSON_파일_저장(s, 파일명)
}

func (s *S내재가치_정보_모음) DB읽기(db *sql.DB) (에러 error) {
	defer lib.S예외처리{M에러: &에러, M함수: func() { s = nil }}.S실행()

	lib.F확인(F내재가치_정보_테이블_생성(db))

	SQL := new(bytes.Buffer)
	SQL.WriteString("SELECT code, date, json ")
	SQL.WriteString("FROM fundamental_data ")
	SQL.WriteString("ORDER BY code, date")

	stmt, 에러 := db.Prepare(SQL.String())
	lib.F확인(에러)
	defer stmt.Close()

	rows, 에러 := stmt.Query()
	lib.F확인(에러)
	defer rows.Close()

	for rows.Next() {
		var code string
		var date time.Time
		var json string
		var 값 *S내재가치_정보

		lib.F확인(rows.Scan(&code, &date, &json))

		if 에러 = lib.F디코딩(lib.JSON, []byte(json), &값); 에러 != nil {
			lib.F문자열_출력("%v %v : 디코딩 에러\n%v", code, date.Format(lib.P일자_형식), 에러)
			continue
		} else if 값 == nil {
			lib.F문자열_출력("%v %v : nil 값", code, date.Format(lib.P일자_형식))
			continue
		} else if 값.S내재가치_식별정보 == nil {
			lib.F문자열_출력("%v %v : nil 식별정보", code, date.Format(lib.P일자_형식))
			continue
		}

		s.M저장소[값.G키()] = 값
	}

	return
}

func (s *S내재가치_정보_모음) DB저장(db *sql.DB) error {
	// 맵 데이터를 슬라이스 형태로 변환
	값_모음 := make([]*S내재가치_정보, len(s.M저장소))

	i := 0
	for _, 값 := range s.M저장소 {
		값_모음[i] = 값
		i++
	}

	// 슬라이스 데이터 DB에 저장
	return F내재가치_정보_모음_DB저장(db, 값_모음)
}

type S내재가치_정보 struct {
	*S내재가치_식별정보
	*S재무제표_정보_내용
	*S재무비율_정보_내용
	*S투자지표_정보_내용
}

type S재무제표_정보 struct {
	*S내재가치_식별정보
	*S재무제표_정보_내용
}

type S재무비율_정보 struct {
	*S내재가치_식별정보
	*S재무비율_정보_내용
}

type S투자지표_정보 struct {
	*S내재가치_식별정보
	*S투자지표_정보_내용
}

func New내재가치_식별정보(종목코드 string, 연도 uint16, 월 uint8) *S내재가치_식별정보 {
	s := new(S내재가치_식별정보)
	s.M종목코드 = 종목코드
	s.S연도_월(연도, 월)

	return s
}

type S내재가치_식별정보 struct {
	M종목코드 string
	M연도_월 uint32
}

func (s S내재가치_식별정보) G키() string {
	return s.M종목코드 + strconv.Itoa(int(s.M연도_월))
}

func (s *S내재가치_식별정보) S연도_월(연도 uint16, 월 uint8) {
	s.M연도_월 = uint32(연도)*100 + uint32(월)
}

func (s S내재가치_식별정보) G일자() uint32 {
	return s.M연도_월*100 + 1
}

func (s S내재가치_식별정보) G일자2() time.Time {
	return lib.F2포맷된_일자_단순형("20060102", strconv.Itoa(int(s.M연도_월*100+1)))
}

type S재무제표_정보_내용 struct {
	M매출액      float64
	M매출총이익    float64
	M영업이익     float64
	M세전계속사업이익 float64
	M당기순이익    float64
	M자산       float64
	M유동자산     float64 // NCAV 전략에 필요함.
	M현금및현금성자산 float64 // EV계산에 필요함.
	M부채       float64
	M유동부채     float64
	M자본       float64
	M영업_현금흐름  float64
	M투자_현금흐름  float64
	M재무_현금흐름  float64
	M현금_증가    float64
}

type S재무비율_정보_내용 struct {
	M유동_비율     float64
	M당좌_비율     float64
	M부채_비율     float64
	M이자보상배율    float64
	M매출총이익율    float64
	M세전계속사업이익율 float64
	M영업이익율     float64
	EBITDA마진율  float64
	ROA        float64
	ROE        float64
	ROIC       float64
	M총자산회전율    float64
	M총부채회전율    float64
	M총자본회전율    float64
	M순운전자본회전율  float64
}

type S투자지표_정보_내용 struct {
	M상장주식수    float64
	EPS       float64
	EBITDAPS  float64
	CFPS      float64
	SPS       float64
	BPS       float64
	DPS       float64
	M배당성향     float64
	M총현금흐름    float64
	M세후영업이익   float64
	M유무형자산상각비 float64
	M총투자      float64
	FCFF      float64
}

func F내재가치_정보_모음_DB저장(db *sql.DB, 값_모음 []*S내재가치_정보) (에러 error) {
	var tx *sql.Tx
	defer lib.S예외처리{M에러: &에러, M함수: func() { lib.F조건부_실행(tx != nil, tx.Rollback) }}.S실행()

	lib.F확인(F내재가치_정보_테이블_생성(db))

	SQL생성 := new(bytes.Buffer)
	SQL생성.WriteString("INSERT IGNORE INTO fundamental_data (")
	SQL생성.WriteString("code,")
	SQL생성.WriteString("date,")
	SQL생성.WriteString("json")
	SQL생성.WriteString(") VALUES (?,?,'')")

	SQL수정 := new(bytes.Buffer)
	SQL수정.WriteString("UPDATE fundamental_data SET json=? ")
	SQL수정.WriteString("WHERE code=? AND date=?")

	txOpts := new(sql.TxOptions)
	txOpts.Isolation = sql.LevelDefault
	txOpts.ReadOnly = false

	tx, 에러 = db.BeginTx(context.TODO(), txOpts)
	lib.F확인(에러)

	stmt생성, 에러 := tx.Prepare(SQL생성.String())
	lib.F확인(에러)
	defer stmt생성.Close()

	stmt수정, 에러 := tx.Prepare(SQL수정.String())
	lib.F확인(에러)
	defer stmt수정.Close()

	for _, 값 := range 값_모음 {
		json, 에러 := lib.F인코딩(lib.JSON, 값)
		lib.F확인(에러)

		_, 에러 = stmt생성.Exec(값.M종목코드, 값.G일자())
		lib.F확인(에러)

		_, 에러 = stmt수정.Exec(string(json), 값.M종목코드, 값.G일자())
		lib.F확인(에러)
	}

	tx.Commit()

	return nil
}

func F내재가치_정보_테이블_생성(db *sql.DB) error {
	SQL := new(bytes.Buffer)
	SQL.WriteString("CREATE TABLE IF NOT EXISTS fundamental_data (")
	SQL.WriteString("code CHAR(8) NOT NULL,")
	SQL.WriteString("date DATE NOT NULL,")
	SQL.WriteString("json TEXT NOT NULL,")
	SQL.WriteString("CONSTRAINT PRIMARY KEY (code,date)")
	SQL.WriteString(")")

	_, 에러 := db.Exec(SQL.String())

	return 에러
}
<|endoftext|>


<commit_before>// Package main implements a simple HTTP interface for zfs management
package zfsd

import (
	"errors"
	"fmt"
	"net/http"

	"gopkg.in/mistifyio/go-zfs.v1"
)

type (

	// ZFS is used for RPC services
	ZFS struct {
	}

	ListRequest struct {
		Type   string `json:"type"`
		Prefix string `json:"prefix"`
	}

	GetRequest struct {
		Name string `json:"name"`
	}

	SetRequest struct {
		Name       string            `json:"name"`
		Properties map[string]string `json:"properties,omitempty"`
	}

	SnapshotRequest struct {
		Name     string `json:"name"`
		Snapshot string `json:"snapshot"`
	}

	RollbackRequest struct {
		Name      string `json:"name"`
		Snapshot  string `json:"snapshot"`
		Recursive bool   `json:"recursive"`
	}

	CloneRequest struct {
		Name       string            `json:"name"`
		Snapshot   string            `json:"snapshot"`
		Target     string            `json:"target"`
		Properties map[string]string `json:"properties,omitempty"`
	}

	DestroyRequest struct {
		Name      string `json:"name"`
		Recursive bool   `json:"recursive"`
	}
)

// List retrieves a list of all ZFS datasets, optionally only of a certain type or prefix
func (z *ZFS) List(r *http.Request, req *ListRequest, resp *[]*zfs.Dataset) error {

	var ds []*zfs.Dataset
	var err error
	switch req.Type {
	case "snapshot":
		ds, err = zfs.Snapshots(req.Prefix)
	case "filesystem":
		ds, err = zfs.Filesystems(req.Prefix)
	case "volume":
		ds, err = zfs.Volumes(req.Prefix)
	case "", "all":
		ds, err = zfs.Datasets(req.Prefix)
	default:
		fmt.Errorf("unknown type: %s", req.Type)
	}
	if err != nil {
		return err
	}
	*resp = ds
	return nil
}

// Get retrieves a single ZFS dataset.
func (z *ZFS) Get(r *http.Request, req *GetRequest, resp *zfs.Dataset) error {
	if req.Name == "" {
		return errors.New("must have name")
	}

	ds, err := zfs.GetDataset(req.Name)
	if err != nil {
		return err
	}
	*resp = *ds
	return nil
}

// Set sets properties on a ZFS dataset.
func (z *ZFS) Set(r *http.Request, req *SetRequest, resp *zfs.Dataset) error {
	if req.Name == "" {
		return errors.New("must have name")
	}
	if req.Properties == nil || len(req.Properties) == 0 {
		return errors.New("must have properties")
	}

	ds, err := zfs.GetDataset(req.Name)
	if err != nil {
		return err
	}

	// zfs should have a setproperties that takes a map
	for k, v := range req.Properties {
		err := ds.SetProperty(k, v)
		if err != nil {
			return err
		}
	}
	ds, err = zfs.GetDataset(req.Name)
	if err != nil {
		return err
	}
	*resp = *ds
	return nil
}

// Snapshot creates a snapshot of a ZFS dataset.
func (z *ZFS) Snapshot(r *http.Request, req *SnapshotRequest, resp *zfs.Dataset) error {
	if req.Name == "" {
		return errors.New("must have name")
	}

	if req.Snapshot == "" {
		return errors.New("must have snapshot")
	}

	ds, err := zfs.GetDataset(req.Name)
	if err != nil {
		return err
	}

	snap, err := ds.Snapshot(req.Snapshot, false)
	if err != nil {
		return err
	}

	*resp = *snap
	return nil
}

// Clone clones a ZFS snapshot
func (z *ZFS) Clone(r *http.Request, req *CloneRequest, resp *zfs.Dataset) error {
	if req.Name == "" {
		return errors.New("must have name")
	}

	if req.Snapshot == "" {
		return errors.New("must have snapshot")
	}

	if req.Target == "" {
		return errors.New("must have target")
	}

	ds, err := zfs.GetDataset(req.Name)
	if err != nil {
		return err
	}

	clone, err := ds.Clone(req.Target, req.Properties)
	if err != nil {
		return err
	}

	*resp = *clone
	return nil
}

// Destroy removes a ZFS dataset
func (z *ZFS) Destroy(r *http.Request, req *DestroyRequest, resp *zfs.Dataset) error {
	if req.Name == "" {
		return errors.New("must have name")
	}

	ds, err := zfs.GetDataset(req.Name)
	if err != nil {
		return err
	}

	err = ds.Destroy(req.Recursive)
	if err != nil {
		return err
	}
	*resp = *ds
	return nil
}

// Rollback rolls the given dataset to back a previous snapshot.
func (z *ZFS) Rollback(r *http.Request, req *RollbackRequest, resp *zfs.Dataset) error {
	if req.Name == "" {
		return errors.New("must have name")
	}

	if req.Snapshot == "" {
		return errors.New("must have snapshot")
	}

	snapName := fmt.Sprintf("%s@%s", req.Name, req.Snapshot)
	snap, err := zfs.GetDataset(snapName)
	if err != nil {
		return err
	}

	err = snap.Rollback(false)
	if err != nil {
		return err
	}

	ds, err := zfs.GetDataset(req.Name)
	if err != nil {
		return err
	}

	*resp = *ds
	return nil
}
<commit_msg>fix snapshot get<commit_after>// Package main implements a simple HTTP interface for zfs management
package zfsd

import (
	"errors"
	"fmt"
	"net/http"

	"gopkg.in/mistifyio/go-zfs.v1"
)

type (

	// ZFS is used for RPC services
	ZFS struct {
	}

	ListRequest struct {
		Type   string `json:"type"`
		Prefix string `json:"prefix"`
	}

	GetRequest struct {
		Name string `json:"name"`
	}

	SetRequest struct {
		Name       string            `json:"name"`
		Properties map[string]string `json:"properties,omitempty"`
	}

	SnapshotRequest struct {
		Name     string `json:"name"`
		Snapshot string `json:"snapshot"`
	}

	RollbackRequest struct {
		Name      string `json:"name"`
		Snapshot  string `json:"snapshot"`
		Recursive bool   `json:"recursive"`
	}

	CloneRequest struct {
		Name       string            `json:"name"`
		Snapshot   string            `json:"snapshot"`
		Target     string            `json:"target"`
		Properties map[string]string `json:"properties,omitempty"`
	}

	DestroyRequest struct {
		Name      string `json:"name"`
		Recursive bool   `json:"recursive"`
	}
)

// List retrieves a list of all ZFS datasets, optionally only of a certain type or prefix
func (z *ZFS) List(r *http.Request, req *ListRequest, resp *[]*zfs.Dataset) error {

	var ds []*zfs.Dataset
	var err error
	switch req.Type {
	case "snapshot":
		ds, err = zfs.Snapshots(req.Prefix)
	case "filesystem":
		ds, err = zfs.Filesystems(req.Prefix)
	case "volume":
		ds, err = zfs.Volumes(req.Prefix)
	case "", "all":
		ds, err = zfs.Datasets(req.Prefix)
	default:
		fmt.Errorf("unknown type: %s", req.Type)
	}
	if err != nil {
		return err
	}
	*resp = ds
	return nil
}

// Get retrieves a single ZFS dataset.
func (z *ZFS) Get(r *http.Request, req *GetRequest, resp *zfs.Dataset) error {
	if req.Name == "" {
		return errors.New("must have name")
	}

	ds, err := zfs.GetDataset(req.Name)
	if err != nil {
		return err
	}
	*resp = *ds
	return nil
}

// Set sets properties on a ZFS dataset.
func (z *ZFS) Set(r *http.Request, req *SetRequest, resp *zfs.Dataset) error {
	if req.Name == "" {
		return errors.New("must have name")
	}
	if req.Properties == nil || len(req.Properties) == 0 {
		return errors.New("must have properties")
	}

	ds, err := zfs.GetDataset(req.Name)
	if err != nil {
		return err
	}

	// zfs should have a setproperties that takes a map
	for k, v := range req.Properties {
		err := ds.SetProperty(k, v)
		if err != nil {
			return err
		}
	}
	ds, err = zfs.GetDataset(req.Name)
	if err != nil {
		return err
	}
	*resp = *ds
	return nil
}

// Snapshot creates a snapshot of a ZFS dataset.
func (z *ZFS) Snapshot(r *http.Request, req *SnapshotRequest, resp *zfs.Dataset) error {
	if req.Name == "" {
		return errors.New("must have name")
	}

	if req.Snapshot == "" {
		return errors.New("must have snapshot")
	}

	ds, err := zfs.GetDataset(req.Name)
	if err != nil {
		return err
	}

	snap, err := ds.Snapshot(req.Snapshot, false)
	if err != nil {
		return err
	}

	*resp = *snap
	return nil
}

// Clone clones a ZFS snapshot
func (z *ZFS) Clone(r *http.Request, req *CloneRequest, resp *zfs.Dataset) error {
	if req.Name == "" {
		return errors.New("must have name")
	}

	if req.Snapshot == "" {
		return errors.New("must have snapshot")
	}

	if req.Target == "" {
		return errors.New("must have target")
	}

	snapName := fmt.Sprintf("%s@%s", req.Name, req.Snapshot)
	snap, err := zfs.GetDataset(snapName)
	if err != nil {
		return err
	}

	clone, err := snap.Clone(req.Target, req.Properties)
	if err != nil {
		return err
	}

	*resp = *clone
	return nil
}

// Destroy removes a ZFS dataset
func (z *ZFS) Destroy(r *http.Request, req *DestroyRequest, resp *zfs.Dataset) error {
	if req.Name == "" {
		return errors.New("must have name")
	}

	ds, err := zfs.GetDataset(req.Name)
	if err != nil {
		return err
	}

	err = ds.Destroy(req.Recursive)
	if err != nil {
		return err
	}
	*resp = *ds
	return nil
}

// Rollback rolls the given dataset to back a previous snapshot.
func (z *ZFS) Rollback(r *http.Request, req *RollbackRequest, resp *zfs.Dataset) error {
	if req.Name == "" {
		return errors.New("must have name")
	}

	if req.Snapshot == "" {
		return errors.New("must have snapshot")
	}

	snapName := fmt.Sprintf("%s@%s", req.Name, req.Snapshot)
	snap, err := zfs.GetDataset(snapName)
	if err != nil {
		return err
	}

	err = snap.Rollback(false)
	if err != nil {
		return err
	}

	ds, err := zfs.GetDataset(req.Name)
	if err != nil {
		return err
	}

	*resp = *ds
	return nil
}
<|endoftext|>


<commit_before>package server

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path"
	"path/filepath"

	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/health"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/peer"

	"cache/cluster"
	pb "cache/proto/rpc_cache"
)

// maxMsgSize is the maximum message size our gRPC server accepts.
// We deliberately set this to something high since we don't want to limit artifact size here.
const maxMsgSize = 200 * 1024 * 1024

func init() {
	// When tracing is enabled, it appears to keep references to messages alive, possibly indefinitely (?).
	// This is very bad for us since our messages are large, it can result in leaking memory very quickly
	// and ultimately OOM errors. Disabling tracing appears to alleviate the problem.
	grpc.EnableTracing = false
}

type RpcCacheServer struct {
	cache        *Cache
	readonlyKeys map[string]*x509.Certificate
	writableKeys map[string]*x509.Certificate
	cluster      *cluster.Cluster
}

func (r *RpcCacheServer) Store(ctx context.Context, req *pb.StoreRequest) (*pb.StoreResponse, error) {
	if err := r.authenticateClient(r.writableKeys, ctx); err != nil {
		return nil, err
	}
	success := storeArtifact(r.cache, req.Os, req.Arch, req.Hash, req.Artifacts)
	if success && r.cluster != nil {
		// Replicate this artifact to another node. Doesn't have to be done synchronously.
		go r.cluster.ReplicateArtifacts(req)
	}
	return &pb.StoreResponse{Success: success}, nil
}

// storeArtifact stores a series of artifacts in the cache.
// Broken out of above to share with Replicate below.
func storeArtifact(cache *Cache, os, arch string, hash []byte, artifacts []*pb.Artifact) bool {
	arch = os + "_" + arch
	hashStr := base64.RawURLEncoding.EncodeToString(hash)
	for _, artifact := range artifacts {
		path := path.Join(arch, artifact.Package, artifact.Target, hashStr, artifact.File)
		if err := cache.StoreArtifact(path, artifact.Body); err != nil {
			return false
		}
	}
	return true
}

func (r *RpcCacheServer) Retrieve(ctx context.Context, req *pb.RetrieveRequest) (*pb.RetrieveResponse, error) {
	if err := r.authenticateClient(r.readonlyKeys, ctx); err != nil {
		return nil, err
	}
	response := pb.RetrieveResponse{Success: true}
	arch := req.Os + "_" + req.Arch
	hash := base64.RawURLEncoding.EncodeToString(req.Hash)
	for _, artifact := range req.Artifacts {
		root := path.Join(arch, artifact.Package, artifact.Target, hash)
		fileRoot := path.Join(root, artifact.File)
		art, err := r.cache.RetrieveArtifact(fileRoot)
		if err != nil {
			log.Debug("Failed to retrieve artifact %s: %s", fileRoot, err)
			return &pb.RetrieveResponse{Success: false}, nil
		}
		for name, body := range art {
			response.Artifacts = append(response.Artifacts, &pb.Artifact{
				Package: artifact.Package,
				Target:  artifact.Target,
				File:    name[len(root)+1:],
				Body:    body,
			})
		}
	}
	return &response, nil
}

func (r *RpcCacheServer) Delete(ctx context.Context, req *pb.DeleteRequest) (*pb.DeleteResponse, error) {
	if err := r.authenticateClient(r.writableKeys, ctx); err != nil {
		return nil, err
	}
	if req.Everything {
		return &pb.DeleteResponse{Success: r.cache.DeleteAllArtifacts() == nil}, nil
	}
	success := deleteArtifact(r.cache, req.Os, req.Arch, req.Artifacts)
	if success && r.cluster != nil {
		// Delete this artifact from other nodes. Doesn't have to be done synchronously.
		go r.cluster.DeleteArtifacts(req)
	}
	return &pb.DeleteResponse{Success: success}, nil
}

// deleteArtifact handles the actual removal of artifacts from the cache.
// It's split out from Delete to share with replication RPCs below.
func deleteArtifact(cache *Cache, os, arch string, artifacts []*pb.Artifact) bool {
	success := true
	for _, artifact := range artifacts {
		if cache.DeleteArtifact(path.Join(os+"_"+arch, artifact.Package, artifact.Target)) != nil {
			success = false
		}
	}
	return success
}

func (r *RpcCacheServer) ListNodes(ctx context.Context, req *pb.ListRequest) (*pb.ListResponse, error) {
	if err := r.authenticateClient(r.readonlyKeys, ctx); err != nil {
		return nil, err
	}
	if r.cluster == nil {
		return &pb.ListResponse{}, nil
	}
	return &pb.ListResponse{Nodes: r.cluster.GetMembers()}, nil
}

func (r *RpcCacheServer) authenticateClient(certs map[string]*x509.Certificate, ctx context.Context) error {
	if len(certs) == 0 {
		return nil // Open to anyone.
	}
	p, ok := peer.FromContext(ctx)
	if !ok {
		return fmt.Errorf("Missing client certificate")
	}
	info, ok := p.AuthInfo.(credentials.TLSInfo)
	if !ok {
		return fmt.Errorf("Could not extract auth info")
	}
	if len(info.State.PeerCertificates) == 0 {
		return fmt.Errorf("No peer certificate available")
	}
	cert := info.State.PeerCertificates[0]
	if okCert := certs[string(cert.RawSubject)]; okCert != nil && okCert.Equal(cert) {
		return nil
	}
	return fmt.Errorf("Invalid or unknown certificate")
}

func loadKeys(filename string) map[string]*x509.Certificate {
	ret := map[string]*x509.Certificate{}
	if err := filepath.Walk(filename, func(name string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		} else if !info.IsDir() {
			data, err := ioutil.ReadFile(filename)
			if err != nil {
				log.Fatalf("Failed to read cert from %s: %s", filename, err)
			}
			p, _ := pem.Decode(data)
			if p == nil {
				log.Fatalf("Couldn't decode PEM data from %s: %s", filename, err)
			}
			cert, err := x509.ParseCertificate(p.Bytes)
			if err != nil {
				log.Fatalf("Couldn't parse certificate from %s: %s", filename, err)
			}
			ret[string(cert.RawSubject)] = cert
		}
		return nil
	}); err != nil {
		log.Fatalf("%s", err)
	}
	return ret
}

// RPCServer implements the gRPC server for communication between cache nodes.
type RPCServer struct {
	cache   *Cache
	cluster *cluster.Cluster
}

func (r *RPCServer) Join(ctx context.Context, req *pb.JoinRequest) (*pb.JoinResponse, error) {
	// TODO(pebers): Authentication.
	return r.cluster.AddNode(req), nil
}

func (r *RPCServer) Replicate(ctx context.Context, req *pb.ReplicateRequest) (*pb.ReplicateResponse, error) {
	// TODO(pebers): Authentication.
	if req.Delete {
		return &pb.ReplicateResponse{
			Success: deleteArtifact(r.cache, req.Os, req.Arch, req.Artifacts),
		}, nil
	}
	return &pb.ReplicateResponse{
		Success: storeArtifact(r.cache, req.Os, req.Arch, req.Hash, req.Artifacts),
	}, nil
}

// BuildGrpcServer creates a new, unstarted grpc.Server and returns it.
// It also returns a net.Listener to start it on.
func BuildGrpcServer(port int, cache *Cache, cluster *cluster.Cluster, keyFile, certFile, caCertFile, readonlyKeys, writableKeys string) (*grpc.Server, net.Listener) {
	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		log.Fatalf("Failed to listen on port %d: %v", port, err)
	}
	s := serverWithAuth(keyFile, certFile, caCertFile)
	r := &RpcCacheServer{cache: cache, cluster: cluster}
	if writableKeys != "" {
		r.writableKeys = loadKeys(writableKeys)
	}
	if readonlyKeys != "" {
		r.readonlyKeys = loadKeys(readonlyKeys)
		if len(r.readonlyKeys) > 0 {
			// This saves duplication when checking later; writable keys are implicitly readable too.
			for k, v := range r.writableKeys {
				if _, present := r.readonlyKeys[k]; !present {
					r.readonlyKeys[k] = v
				}
			}
		}
	}
	r2 := &RPCServer{cache: cache, cluster: cluster}
	pb.RegisterRpcCacheServer(s, r)
	pb.RegisterRpcServerServer(s, r2)
	healthserver := health.NewServer()
	healthserver.SetServingStatus("plz-rpc-cache", healthpb.HealthCheckResponse_SERVING)
	healthpb.RegisterHealthServer(s, healthserver)
	return s, lis
}

// ServeGrpcForever serves gRPC until killed using the given server.
// It's very simple and provided as a convenience so callers don't have to import grpc themselves.
func ServeGrpcForever(server *grpc.Server, lis net.Listener) {
	log.Notice("Serving RPC cache on %s", lis.Addr())
	server.Serve(lis)
}

// serverWithAuth builds a gRPC server, possibly with authentication if key / cert files are given.
func serverWithAuth(keyFile, certFile, caCertFile string) *grpc.Server {
	if keyFile == "" {
		return grpc.NewServer(grpc.MaxMsgSize(maxMsgSize)) // No auth.
	}
	log.Debug("Loading x509 key pair from key: %s cert: %s", keyFile, certFile)
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		log.Fatalf("Failed to load x509 key pair: %s", err)
	}
	config := tls.Config{
		Certificates: []tls.Certificate{cert},
		ClientAuth:   tls.RequestClientCert,
	}
	if caCertFile != "" {
		cert, err := ioutil.ReadFile(caCertFile)
		if err != nil {
			log.Fatalf("Failed to read CA cert file: %s", err)
		}
		config.ClientCAs = x509.NewCertPool()
		if !config.ClientCAs.AppendCertsFromPEM(cert) {
			log.Fatalf("Failed to find any PEM certificates in CA cert")
		}
	}
	return grpc.NewServer(grpc.Creds(credentials.NewTLS(&config)), grpc.MaxMsgSize(maxMsgSize))
}
<commit_msg>fix rpc_server which can't load certificates from a directory (#222)<commit_after>package server

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path"
	"path/filepath"

	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/health"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/peer"

	"cache/cluster"
	pb "cache/proto/rpc_cache"
)

// maxMsgSize is the maximum message size our gRPC server accepts.
// We deliberately set this to something high since we don't want to limit artifact size here.
const maxMsgSize = 200 * 1024 * 1024

func init() {
	// When tracing is enabled, it appears to keep references to messages alive, possibly indefinitely (?).
	// This is very bad for us since our messages are large, it can result in leaking memory very quickly
	// and ultimately OOM errors. Disabling tracing appears to alleviate the problem.
	grpc.EnableTracing = false
}

type RpcCacheServer struct {
	cache        *Cache
	readonlyKeys map[string]*x509.Certificate
	writableKeys map[string]*x509.Certificate
	cluster      *cluster.Cluster
}

func (r *RpcCacheServer) Store(ctx context.Context, req *pb.StoreRequest) (*pb.StoreResponse, error) {
	if err := r.authenticateClient(r.writableKeys, ctx); err != nil {
		return nil, err
	}
	success := storeArtifact(r.cache, req.Os, req.Arch, req.Hash, req.Artifacts)
	if success && r.cluster != nil {
		// Replicate this artifact to another node. Doesn't have to be done synchronously.
		go r.cluster.ReplicateArtifacts(req)
	}
	return &pb.StoreResponse{Success: success}, nil
}

// storeArtifact stores a series of artifacts in the cache.
// Broken out of above to share with Replicate below.
func storeArtifact(cache *Cache, os, arch string, hash []byte, artifacts []*pb.Artifact) bool {
	arch = os + "_" + arch
	hashStr := base64.RawURLEncoding.EncodeToString(hash)
	for _, artifact := range artifacts {
		path := path.Join(arch, artifact.Package, artifact.Target, hashStr, artifact.File)
		if err := cache.StoreArtifact(path, artifact.Body); err != nil {
			return false
		}
	}
	return true
}

func (r *RpcCacheServer) Retrieve(ctx context.Context, req *pb.RetrieveRequest) (*pb.RetrieveResponse, error) {
	if err := r.authenticateClient(r.readonlyKeys, ctx); err != nil {
		return nil, err
	}
	response := pb.RetrieveResponse{Success: true}
	arch := req.Os + "_" + req.Arch
	hash := base64.RawURLEncoding.EncodeToString(req.Hash)
	for _, artifact := range req.Artifacts {
		root := path.Join(arch, artifact.Package, artifact.Target, hash)
		fileRoot := path.Join(root, artifact.File)
		art, err := r.cache.RetrieveArtifact(fileRoot)
		if err != nil {
			log.Debug("Failed to retrieve artifact %s: %s", fileRoot, err)
			return &pb.RetrieveResponse{Success: false}, nil
		}
		for name, body := range art {
			response.Artifacts = append(response.Artifacts, &pb.Artifact{
				Package: artifact.Package,
				Target:  artifact.Target,
				File:    name[len(root)+1:],
				Body:    body,
			})
		}
	}
	return &response, nil
}

func (r *RpcCacheServer) Delete(ctx context.Context, req *pb.DeleteRequest) (*pb.DeleteResponse, error) {
	if err := r.authenticateClient(r.writableKeys, ctx); err != nil {
		return nil, err
	}
	if req.Everything {
		return &pb.DeleteResponse{Success: r.cache.DeleteAllArtifacts() == nil}, nil
	}
	success := deleteArtifact(r.cache, req.Os, req.Arch, req.Artifacts)
	if success && r.cluster != nil {
		// Delete this artifact from other nodes. Doesn't have to be done synchronously.
		go r.cluster.DeleteArtifacts(req)
	}
	return &pb.DeleteResponse{Success: success}, nil
}

// deleteArtifact handles the actual removal of artifacts from the cache.
// It's split out from Delete to share with replication RPCs below.
func deleteArtifact(cache *Cache, os, arch string, artifacts []*pb.Artifact) bool {
	success := true
	for _, artifact := range artifacts {
		if cache.DeleteArtifact(path.Join(os+"_"+arch, artifact.Package, artifact.Target)) != nil {
			success = false
		}
	}
	return success
}

func (r *RpcCacheServer) ListNodes(ctx context.Context, req *pb.ListRequest) (*pb.ListResponse, error) {
	if err := r.authenticateClient(r.readonlyKeys, ctx); err != nil {
		return nil, err
	}
	if r.cluster == nil {
		return &pb.ListResponse{}, nil
	}
	return &pb.ListResponse{Nodes: r.cluster.GetMembers()}, nil
}

func (r *RpcCacheServer) authenticateClient(certs map[string]*x509.Certificate, ctx context.Context) error {
	if len(certs) == 0 {
		return nil // Open to anyone.
	}
	p, ok := peer.FromContext(ctx)
	if !ok {
		return fmt.Errorf("Missing client certificate")
	}
	info, ok := p.AuthInfo.(credentials.TLSInfo)
	if !ok {
		return fmt.Errorf("Could not extract auth info")
	}
	if len(info.State.PeerCertificates) == 0 {
		return fmt.Errorf("No peer certificate available")
	}
	cert := info.State.PeerCertificates[0]
	if okCert := certs[string(cert.RawSubject)]; okCert != nil && okCert.Equal(cert) {
		return nil
	}
	return fmt.Errorf("Invalid or unknown certificate")
}

func loadKeys(filename string) map[string]*x509.Certificate {
	ret := map[string]*x509.Certificate{}
	if err := filepath.Walk(filename, func(name string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		} else if !info.IsDir() {
			data, err := ioutil.ReadFile(name)
			if err != nil {
				log.Fatalf("Failed to read cert from %s: %s", name, err)
			}
			p, _ := pem.Decode(data)
			if p == nil {
				log.Fatalf("Couldn't decode PEM data from %s: %s", name, err)
			}
			cert, err := x509.ParseCertificate(p.Bytes)
			if err != nil {
				log.Fatalf("Couldn't parse certificate from %s: %s", name, err)
			}
			ret[string(cert.RawSubject)] = cert
		}
		return nil
	}); err != nil {
		log.Fatalf("%s", err)
	}
	return ret
}

// RPCServer implements the gRPC server for communication between cache nodes.
type RPCServer struct {
	cache   *Cache
	cluster *cluster.Cluster
}

func (r *RPCServer) Join(ctx context.Context, req *pb.JoinRequest) (*pb.JoinResponse, error) {
	// TODO(pebers): Authentication.
	return r.cluster.AddNode(req), nil
}

func (r *RPCServer) Replicate(ctx context.Context, req *pb.ReplicateRequest) (*pb.ReplicateResponse, error) {
	// TODO(pebers): Authentication.
	if req.Delete {
		return &pb.ReplicateResponse{
			Success: deleteArtifact(r.cache, req.Os, req.Arch, req.Artifacts),
		}, nil
	}
	return &pb.ReplicateResponse{
		Success: storeArtifact(r.cache, req.Os, req.Arch, req.Hash, req.Artifacts),
	}, nil
}

// BuildGrpcServer creates a new, unstarted grpc.Server and returns it.
// It also returns a net.Listener to start it on.
func BuildGrpcServer(port int, cache *Cache, cluster *cluster.Cluster, keyFile, certFile, caCertFile, readonlyKeys, writableKeys string) (*grpc.Server, net.Listener) {
	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		log.Fatalf("Failed to listen on port %d: %v", port, err)
	}
	s := serverWithAuth(keyFile, certFile, caCertFile)
	r := &RpcCacheServer{cache: cache, cluster: cluster}
	if writableKeys != "" {
		r.writableKeys = loadKeys(writableKeys)
	}
	if readonlyKeys != "" {
		r.readonlyKeys = loadKeys(readonlyKeys)
		if len(r.readonlyKeys) > 0 {
			// This saves duplication when checking later; writable keys are implicitly readable too.
			for k, v := range r.writableKeys {
				if _, present := r.readonlyKeys[k]; !present {
					r.readonlyKeys[k] = v
				}
			}
		}
	}
	r2 := &RPCServer{cache: cache, cluster: cluster}
	pb.RegisterRpcCacheServer(s, r)
	pb.RegisterRpcServerServer(s, r2)
	healthserver := health.NewServer()
	healthserver.SetServingStatus("plz-rpc-cache", healthpb.HealthCheckResponse_SERVING)
	healthpb.RegisterHealthServer(s, healthserver)
	return s, lis
}

// ServeGrpcForever serves gRPC until killed using the given server.
// It's very simple and provided as a convenience so callers don't have to import grpc themselves.
func ServeGrpcForever(server *grpc.Server, lis net.Listener) {
	log.Notice("Serving RPC cache on %s", lis.Addr())
	server.Serve(lis)
}

// serverWithAuth builds a gRPC server, possibly with authentication if key / cert files are given.
func serverWithAuth(keyFile, certFile, caCertFile string) *grpc.Server {
	if keyFile == "" {
		return grpc.NewServer(grpc.MaxMsgSize(maxMsgSize)) // No auth.
	}
	log.Debug("Loading x509 key pair from key: %s cert: %s", keyFile, certFile)
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		log.Fatalf("Failed to load x509 key pair: %s", err)
	}
	config := tls.Config{
		Certificates: []tls.Certificate{cert},
		ClientAuth:   tls.RequestClientCert,
	}
	if caCertFile != "" {
		cert, err := ioutil.ReadFile(caCertFile)
		if err != nil {
			log.Fatalf("Failed to read CA cert file: %s", err)
		}
		config.ClientCAs = x509.NewCertPool()
		if !config.ClientCAs.AppendCertsFromPEM(cert) {
			log.Fatalf("Failed to find any PEM certificates in CA cert")
		}
	}
	return grpc.NewServer(grpc.Creds(credentials.NewTLS(&config)), grpc.MaxMsgSize(maxMsgSize))
}
<|endoftext|>


<commit_before>package nsf

type Apu struct {
	S1, S2 Square

	Odd        bool
	FC         byte
	FT         byte
	IrqDisable bool
}

type Square struct {
	Envelope
	Timer
	Length
	Sweep
	Duty

	Enable bool
}

type Duty struct {
	Type    byte
	Counter byte
}

type Sweep struct {
	Shift     byte
	Negate    bool
	Period    byte
	Enable    bool
	Divider   byte
	Reset     bool
	NegOffset int
}

type Envelope struct {
	Volume   byte
	Divider  byte
	Counter  byte
	Loop     bool
	Constant bool
	Start    bool
}

type Timer struct {
	Tick   uint16
	Length uint16
}

type Length struct {
	Halt    bool
	Counter byte
}

func (a *Apu) Init() {
	a.S1.Sweep.NegOffset = -1
	for i := uint16(0x4000); i <= 0x400f; i++ {
		a.Write(i, 0)
	}
	a.Write(0x4010, 0x10)
	a.Write(0x4011, 0)
	a.Write(0x4012, 0)
	a.Write(0x4013, 0)
	a.Write(0x4015, 0xf)
	a.Write(0x4017, 0)
}

func (a *Apu) Write(v uint16, b byte) {
	switch v & 0xff {
	case 0x00:
		a.S1.Control1(b)
	case 0x01:
		a.S1.Control2(b)
	case 0x02:
		a.S1.Control3(b)
	case 0x03:
		a.S1.Control4(b)
	case 0x04:
		a.S2.Control1(b)
	case 0x05:
		a.S2.Control2(b)
	case 0x06:
		a.S2.Control3(b)
	case 0x07:
		a.S2.Control4(b)
	case 0x15:
		a.S1.Disable(b&0x1 == 0)
		a.S2.Disable(b&0x2 == 0)
	case 0x17:
		a.FT = 0
		if b&0x80 != 0 {
			a.FC = 5
			a.FrameStep()
		} else {
			a.FC = 4
		}
		a.IrqDisable = b&0x40 != 0
	}
}

func (s *Square) Control1(b byte) {
	s.Envelope.Control(b)
	s.Duty.Control(b)
	s.Length.Halt = b&0x20 != 0
}

func (s *Square) Control2(b byte) {
	s.Sweep.Control(b)
}

func (s *Square) Control3(b byte) {
	s.Timer.Length &= 0xff00
	s.Timer.Length |= uint16(b)
}

func (s *Square) Control4(b byte) {
	s.Timer.Length &= 0xff
	s.Timer.Length |= uint16(b&0x7) << 8
	s.Length.Set(b >> 3)

	s.Envelope.Start = true
	s.Duty.Counter = 0
}

func (d *Duty) Control(b byte) {
	d.Type = b >> 6
}

func (s *Sweep) Control(b byte) {
	s.Shift = b & 0x7
	s.Negate = b&0x8 != 0
	s.Period = (b >> 4) & 0x7
	s.Enable = b&0x80 != 0
	s.Reset = true
}

func (e *Envelope) Control(b byte) {
	e.Volume = b & 0xf
	e.Constant = b&0x10 != 0
	e.Loop = b&0x20 != 0
}

func (l *Length) Set(b byte) {
	if !l.Halt {
		l.Counter = LenLookup[b]
	}
}

func (l *Length) Enabled() bool {
	return l.Counter != 0
}

func (s *Square) Disable(b bool) {
	s.Enable = !b
	if b {
		s.Length.Counter = 0
	}
}

func (a *Apu) Read(v uint16) byte {
	var b byte
	if v == 0x4015 {
		if a.S1.Length.Counter > 0 {
			b |= 0x1
		}
		if a.S2.Length.Counter > 0 {
			b |= 0x2
		}
	}
	return b
}

func (d *Duty) Clock() {
	if d.Counter == 0 {
		d.Counter = 7
	} else {
		d.Counter--
	}
}

func (s *Sweep) Clock() (r bool) {
	if s.Divider == 0 {
		s.Divider = s.Period
		r = true
	} else {
		s.Divider--
	}
	if s.Reset {
		s.Divider = 0
		s.Reset = false
	}
	return
}

func (e *Envelope) Clock() {
	if e.Start {
		e.Start = false
		e.Counter = 15
	} else {
		if e.Divider == 0 {
			e.Divider = e.Volume
			if e.Counter != 0 {
				e.Counter--
			} else if e.Loop {
				e.Counter = 15
			}
		} else {
			e.Divider--
		}
	}
}

// 1.79 MHz/(N+1)
// square: -> duty cycle generator
// triangle: -> triangle step generator
// noise: -> random number generator
func (t *Timer) Clock() bool {
	if t.Tick == 0 {
		t.Tick = t.Length
	} else {
		t.Tick--
	}
	return t.Tick == t.Length
}

func (s *Square) Clock() {
	if s.Timer.Clock() {
		s.Duty.Clock()
	}
}

func (a *Apu) Step() {
	if a.Odd {
		if a.S1.Enable {
			a.S1.Clock()
		}
		if a.S2.Enable {
			a.S2.Clock()
		}
	}
	a.Odd = !a.Odd
	}
}

func (a *Apu) FrameStep() {
	a.FT++
	if a.FT == a.FC {
		a.FT = 0
	}
	if a.FT <= 3 {
		a.S1.Envelope.Clock()
	}
	if a.FT == 1 || a.FT == 3 {
		a.S1.FrameStep()
		a.S2.FrameStep()
	}
	if a.FC == 4 && a.FT == 3 && !a.IrqDisable {
		// todo: assert cpu irq line
	}
}

func (s *Square) FrameStep() {
	s.Length.Clock()
	if s.Sweep.Clock() && s.Sweep.Enable && s.Sweep.Shift > 0 {
		r := s.SweepResult()
		if r <= 0x7ff {
			s.Timer.Tick = r
		}
	}
}

func (l *Length) Clock() {
	if !l.Halt && l.Counter > 0 {
		l.Counter--
	}
}

func (a *Apu) Volume() float32 {
	p := PulseOut[a.S1.Volume()+a.S2.Volume()]
	return p
}

func (s *Square) Volume() uint8 {
	if s.Enable && s.Duty.Enabled() && s.Length.Enabled() && s.Timer.Tick >= 8 && s.SweepResult() <= 0x7ff {
		return s.Envelope.Output()
	}
	return 0
}

func (e *Envelope) Output() byte {
	if e.Constant {
		return e.Volume
	}
	return e.Counter
}

func (s *Square) SweepResult() uint16 {
	r := int(s.Timer.Tick >> s.Sweep.Shift)
	if s.Sweep.Negate {
		r =  - r
	}
	r += int(s.Timer.Tick)
	if r > 0x7ff {
		r = 0x800
	}
	return uint16(r)
}

func (d *Duty) Enabled() bool {
	return DutyCycle[d.Type][d.Counter] == 1
}

var (
	PulseOut  [32]float32
	DutyCycle = [4][8]byte{
		{0, 1, 0, 0, 0, 0, 0, 0},
		{0, 1, 1, 0, 0, 0, 0, 0},
		{0, 1, 1, 1, 1, 0, 0, 0},
		{1, 0, 0, 1, 1, 1, 1, 1},
	}
	LenLookup = []byte{
		0x0a, 0xfe, 0x14, 0x02,
		0x28, 0x04, 0x50, 0x06,
		0xa0, 0x08, 0x3c, 0x0a,
		0x0e, 0x0c, 0x1a, 0x0e,
		0x0c, 0x10, 0x18, 0x12,
		0x30, 0x14, 0x60, 0x16,
		0xc0, 0x18, 0x48, 0x1a,
		0x10, 0x1c, 0x20, 0x1e,
	}
)

func init() {
	for i := range PulseOut {
		PulseOut[i] = 95.88 / (8128/float32(i) + 100)
	}
}
<commit_msg>Remove comment<commit_after>package nsf

type Apu struct {
	S1, S2 Square

	Odd        bool
	FC         byte
	FT         byte
	IrqDisable bool
}

type Square struct {
	Envelope
	Timer
	Length
	Sweep
	Duty

	Enable bool
}

type Duty struct {
	Type    byte
	Counter byte
}

type Sweep struct {
	Shift     byte
	Negate    bool
	Period    byte
	Enable    bool
	Divider   byte
	Reset     bool
	NegOffset int
}

type Envelope struct {
	Volume   byte
	Divider  byte
	Counter  byte
	Loop     bool
	Constant bool
	Start    bool
}

type Timer struct {
	Tick   uint16
	Length uint16
}

type Length struct {
	Halt    bool
	Counter byte
}

func (a *Apu) Init() {
	a.S1.Sweep.NegOffset = -1
	for i := uint16(0x4000); i <= 0x400f; i++ {
		a.Write(i, 0)
	}
	a.Write(0x4010, 0x10)
	a.Write(0x4011, 0)
	a.Write(0x4012, 0)
	a.Write(0x4013, 0)
	a.Write(0x4015, 0xf)
	a.Write(0x4017, 0)
}

func (a *Apu) Write(v uint16, b byte) {
	switch v & 0xff {
	case 0x00:
		a.S1.Control1(b)
	case 0x01:
		a.S1.Control2(b)
	case 0x02:
		a.S1.Control3(b)
	case 0x03:
		a.S1.Control4(b)
	case 0x04:
		a.S2.Control1(b)
	case 0x05:
		a.S2.Control2(b)
	case 0x06:
		a.S2.Control3(b)
	case 0x07:
		a.S2.Control4(b)
	case 0x15:
		a.S1.Disable(b&0x1 == 0)
		a.S2.Disable(b&0x2 == 0)
	case 0x17:
		a.FT = 0
		if b&0x80 != 0 {
			a.FC = 5
			a.FrameStep()
		} else {
			a.FC = 4
		}
		a.IrqDisable = b&0x40 != 0
	}
}

func (s *Square) Control1(b byte) {
	s.Envelope.Control(b)
	s.Duty.Control(b)
	s.Length.Halt = b&0x20 != 0
}

func (s *Square) Control2(b byte) {
	s.Sweep.Control(b)
}

func (s *Square) Control3(b byte) {
	s.Timer.Length &= 0xff00
	s.Timer.Length |= uint16(b)
}

func (s *Square) Control4(b byte) {
	s.Timer.Length &= 0xff
	s.Timer.Length |= uint16(b&0x7) << 8
	s.Length.Set(b >> 3)

	s.Envelope.Start = true
	s.Duty.Counter = 0
}

func (d *Duty) Control(b byte) {
	d.Type = b >> 6
}

func (s *Sweep) Control(b byte) {
	s.Shift = b & 0x7
	s.Negate = b&0x8 != 0
	s.Period = (b >> 4) & 0x7
	s.Enable = b&0x80 != 0
	s.Reset = true
}

func (e *Envelope) Control(b byte) {
	e.Volume = b & 0xf
	e.Constant = b&0x10 != 0
	e.Loop = b&0x20 != 0
}

func (l *Length) Set(b byte) {
	if !l.Halt {
		l.Counter = LenLookup[b]
	}
}

func (l *Length) Enabled() bool {
	return l.Counter != 0
}

func (s *Square) Disable(b bool) {
	s.Enable = !b
	if b {
		s.Length.Counter = 0
	}
}

func (a *Apu) Read(v uint16) byte {
	var b byte
	if v == 0x4015 {
		if a.S1.Length.Counter > 0 {
			b |= 0x1
		}
		if a.S2.Length.Counter > 0 {
			b |= 0x2
		}
	}
	return b
}

func (d *Duty) Clock() {
	if d.Counter == 0 {
		d.Counter = 7
	} else {
		d.Counter--
	}
}

func (s *Sweep) Clock() (r bool) {
	if s.Divider == 0 {
		s.Divider = s.Period
		r = true
	} else {
		s.Divider--
	}
	if s.Reset {
		s.Divider = 0
		s.Reset = false
	}
	return
}

func (e *Envelope) Clock() {
	if e.Start {
		e.Start = false
		e.Counter = 15
	} else {
		if e.Divider == 0 {
			e.Divider = e.Volume
			if e.Counter != 0 {
				e.Counter--
			} else if e.Loop {
				e.Counter = 15
			}
		} else {
			e.Divider--
		}
	}
}

func (t *Timer) Clock() bool {
	if t.Tick == 0 {
		t.Tick = t.Length
	} else {
		t.Tick--
	}
	return t.Tick == t.Length
}

func (s *Square) Clock() {
	if s.Timer.Clock() {
		s.Duty.Clock()
	}
}

func (a *Apu) Step() {
	if a.Odd {
		if a.S1.Enable {
			a.S1.Clock()
		}
		if a.S2.Enable {
			a.S2.Clock()
		}
	}
	a.Odd = !a.Odd
	}
}

func (a *Apu) FrameStep() {
	a.FT++
	if a.FT == a.FC {
		a.FT = 0
	}
	if a.FT <= 3 {
		a.S1.Envelope.Clock()
	}
	if a.FT == 1 || a.FT == 3 {
		a.S1.FrameStep()
		a.S2.FrameStep()
	}
	if a.FC == 4 && a.FT == 3 && !a.IrqDisable {
		// todo: assert cpu irq line
	}
}

func (s *Square) FrameStep() {
	s.Length.Clock()
	if s.Sweep.Clock() && s.Sweep.Enable && s.Sweep.Shift > 0 {
		r := s.SweepResult()
		if r <= 0x7ff {
			s.Timer.Tick = r
		}
	}
}

func (l *Length) Clock() {
	if !l.Halt && l.Counter > 0 {
		l.Counter--
	}
}

func (a *Apu) Volume() float32 {
	p := PulseOut[a.S1.Volume()+a.S2.Volume()]
	return p
}

func (s *Square) Volume() uint8 {
	if s.Enable && s.Duty.Enabled() && s.Length.Enabled() && s.Timer.Tick >= 8 && s.SweepResult() <= 0x7ff {
		return s.Envelope.Output()
	}
	return 0
}

func (e *Envelope) Output() byte {
	if e.Constant {
		return e.Volume
	}
	return e.Counter
}

func (s *Square) SweepResult() uint16 {
	r := int(s.Timer.Tick >> s.Sweep.Shift)
	if s.Sweep.Negate {
		r =  - r
	}
	r += int(s.Timer.Tick)
	if r > 0x7ff {
		r = 0x800
	}
	return uint16(r)
}

func (d *Duty) Enabled() bool {
	return DutyCycle[d.Type][d.Counter] == 1
}

var (
	PulseOut  [32]float32
	DutyCycle = [4][8]byte{
		{0, 1, 0, 0, 0, 0, 0, 0},
		{0, 1, 1, 0, 0, 0, 0, 0},
		{0, 1, 1, 1, 1, 0, 0, 0},
		{1, 0, 0, 1, 1, 1, 1, 1},
	}
	LenLookup = []byte{
		0x0a, 0xfe, 0x14, 0x02,
		0x28, 0x04, 0x50, 0x06,
		0xa0, 0x08, 0x3c, 0x0a,
		0x0e, 0x0c, 0x1a, 0x0e,
		0x0c, 0x10, 0x18, 0x12,
		0x30, 0x14, 0x60, 0x16,
		0xc0, 0x18, 0x48, 0x1a,
		0x10, 0x1c, 0x20, 0x1e,
	}
)

func init() {
	for i := range PulseOut {
		PulseOut[i] = 95.88 / (8128/float32(i) + 100)
	}
}
<|endoftext|>