<commit_before>// Copyright 2018 CNI authors
//
// 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 main

import (
	"encoding/json"
	"fmt"
	"net"
	"strings"

	"github.com/containernetworking/cni/pkg/skel"
	"github.com/containernetworking/cni/pkg/types"
	"github.com/containernetworking/cni/pkg/types/current"
	"github.com/containernetworking/cni/pkg/version"

	types020 "github.com/containernetworking/cni/pkg/types/020"
)

// The top-level network config - IPAM plugins are passed the full configuration
// of the calling plugin, not just the IPAM section.
type Net struct {
	Name       string      `json:"name"`
	CNIVersion string      `json:"cniVersion"`
	IPAM       *IPAMConfig `json:"ipam"`
}

type IPAMConfig struct {
	Name      string
	Type      string         `json:"type"`
	Routes    []*types.Route `json:"routes"`
	Addresses []Address      `json:"addresses,omitempty"`
	DNS       types.DNS      `json:"dns"`
}

type IPAMEnvArgs struct {
	types.CommonArgs
	IP      types.UnmarshallableString `json:"ip,omitempty"`
	GATEWAY types.UnmarshallableString `json:"gateway,omitempty"`
}

type Address struct {
	AddressStr string `json:"address"`
	Gateway    net.IP `json:"gateway,omitempty"`
	Address    net.IPNet
	Version    string
}

func main() {
	// TODO: implement plugin version
	skel.PluginMain(cmdAdd, cmdGet, cmdDel, version.All, "TODO")
}

func cmdGet(args *skel.CmdArgs) error {
	// TODO: implement
	return fmt.Errorf("not implemented")
}

// canonicalizeIP makes sure a provided ip is in standard form
func canonicalizeIP(ip *net.IP) error {
	if ip.To4() != nil {
		*ip = ip.To4()
		return nil
	} else if ip.To16() != nil {
		*ip = ip.To16()
		return nil
	}
	return fmt.Errorf("IP %s not v4 nor v6", *ip)
}

// LoadIPAMConfig creates IPAMConfig using json encoded configuration provided
// as `bytes`. At the moment values provided in envArgs are ignored so there
// is no possibility to overload the json configuration using envArgs
func LoadIPAMConfig(bytes []byte, envArgs string) (*IPAMConfig, string, error) {
	n := Net{}
	if err := json.Unmarshal(bytes, &n); err != nil {
		return nil, "", err
	}

	if n.IPAM == nil {
		return nil, "", fmt.Errorf("IPAM config missing 'ipam' key")
	}

	// Validate all ranges
	numV4 := 0
	numV6 := 0

	for i := range n.IPAM.Addresses {
		ip, addr, err := net.ParseCIDR(n.IPAM.Addresses[i].AddressStr)
		if err != nil {
			return nil, "", fmt.Errorf("invalid CIDR %s: %s", n.IPAM.Addresses[i].AddressStr, err)
		}
		n.IPAM.Addresses[i].Address = *addr
		n.IPAM.Addresses[i].Address.IP = ip

		if err := canonicalizeIP(&n.IPAM.Addresses[i].Address.IP); err != nil {
			return nil, "", fmt.Errorf("invalid address %d: %s", i, err)
		}

		if n.IPAM.Addresses[i].Address.IP.To4() != nil {
			n.IPAM.Addresses[i].Version = "4"
			numV4++
		} else {
			n.IPAM.Addresses[i].Version = "6"
			numV6++
		}
	}

	if envArgs != "" {
		e := IPAMEnvArgs{}
		err := types.LoadArgs(envArgs, &e)
		if err != nil {
			return nil, "", err
		}

		if e.IP != "" {
			for _, item := range strings.Split(string(e.IP), ",") {
				ipstr := strings.TrimSpace(item)

				ip, subnet, err := net.ParseCIDR(ipstr)
				if err != nil {
					return nil, "", fmt.Errorf("invalid CIDR %s: %s", e.IP, err)
				}

				addr := Address{Address: net.IPNet{IP: ip, Mask: subnet.Mask}}
				if addr.Address.IP.To4() != nil {
					addr.Version = "4"
					numV4++
				} else {
					addr.Version = "6"
					numV6++
				}
				n.IPAM.Addresses = append(n.IPAM.Addresses, addr)
			}
		}

		if e.GATEWAY != "" {
			for _, item := range strings.Split(string(e.GATEWAY), ",") {
				gwip := net.ParseIP(strings.TrimSpace(item))
				if gwip == nil {
					return nil, "", fmt.Errorf("invalid gateway address: %s", item)
				}

				for i := range n.IPAM.Addresses {
					if n.IPAM.Addresses[i].Address.Contains(gwip) {
						n.IPAM.Addresses[i].Gateway = gwip
					}
				}
			}
		}
	}

	// CNI spec 0.2.0 and below supported only one v4 and v6 address
	if numV4 > 1 || numV6 > 1 {
		for _, v := range types020.SupportedVersions {
			if n.CNIVersion == v {
				return nil, "", fmt.Errorf("CNI version %v does not support more than 1 address per family", n.CNIVersion)
			}
		}
	}

	// Copy net name into IPAM so not to drag Net struct around
	n.IPAM.Name = n.Name

	return n.IPAM, n.CNIVersion, nil
}

func cmdAdd(args *skel.CmdArgs) error {
	ipamConf, confVersion, err := LoadIPAMConfig(args.StdinData, args.Args)
	if err != nil {
		return err
	}

	result := &current.Result{}
	result.DNS = ipamConf.DNS
	result.Routes = ipamConf.Routes
	for _, v := range ipamConf.Addresses {
		result.IPs = append(result.IPs, &current.IPConfig{
			Version: v.Version,
			Address: v.Address,
			Gateway: v.Gateway})
	}

	return types.PrintResult(result, confVersion)
}

func cmdDel(args *skel.CmdArgs) error {
	// Nothing required because of no resource allocation in static plugin.
	return nil
}
<commit_msg>Incorporate Dan's comments.<commit_after>// Copyright 2018 CNI authors
//
// 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 main

import (
	"encoding/json"
	"fmt"
	"net"
	"strings"

	"github.com/containernetworking/cni/pkg/skel"
	"github.com/containernetworking/cni/pkg/types"
	"github.com/containernetworking/cni/pkg/types/current"
	"github.com/containernetworking/cni/pkg/version"

	types020 "github.com/containernetworking/cni/pkg/types/020"
)

// The top-level network config - IPAM plugins are passed the full configuration
// of the calling plugin, not just the IPAM section.
type Net struct {
	Name       string      `json:"name"`
	CNIVersion string      `json:"cniVersion"`
	IPAM       *IPAMConfig `json:"ipam"`
}

type IPAMConfig struct {
	Name      string
	Type      string         `json:"type"`
	Routes    []*types.Route `json:"routes"`
	Addresses []Address      `json:"addresses,omitempty"`
	DNS       types.DNS      `json:"dns"`
}

type IPAMEnvArgs struct {
	types.CommonArgs
	IP      types.UnmarshallableString `json:"ip,omitempty"`
	GATEWAY types.UnmarshallableString `json:"gateway,omitempty"`
}

type Address struct {
	AddressStr string `json:"address"`
	Gateway    net.IP `json:"gateway,omitempty"`
	Address    net.IPNet
	Version    string
}

func main() {
	// TODO: implement plugin version
	skel.PluginMain(cmdAdd, cmdGet, cmdDel, version.All, "TODO")
}

func cmdGet(args *skel.CmdArgs) error {
	// TODO: implement
	return fmt.Errorf("not implemented")
}

// canonicalizeIP makes sure a provided ip is in standard form
func canonicalizeIP(ip *net.IP) error {
	if ip.To4() != nil {
		*ip = ip.To4()
		return nil
	} else if ip.To16() != nil {
		*ip = ip.To16()
		return nil
	}
	return fmt.Errorf("IP %s not v4 nor v6", *ip)
}

// LoadIPAMConfig creates IPAMConfig using json encoded configuration provided
// as `bytes`. At the moment values provided in envArgs are ignored so there
// is no possibility to overload the json configuration using envArgs
func LoadIPAMConfig(bytes []byte, envArgs string) (*IPAMConfig, string, error) {
	n := Net{}
	if err := json.Unmarshal(bytes, &n); err != nil {
		return nil, "", err
	}

	if n.IPAM == nil {
		return nil, "", fmt.Errorf("IPAM config missing 'ipam' key")
	}

	// Validate all ranges
	numV4 := 0
	numV6 := 0

	for i := range n.IPAM.Addresses {
		ip, addr, err := net.ParseCIDR(n.IPAM.Addresses[i].AddressStr)
		if err != nil {
			return nil, "", fmt.Errorf("invalid CIDR %s: %s", n.IPAM.Addresses[i].AddressStr, err)
		}
		n.IPAM.Addresses[i].Address = *addr
		n.IPAM.Addresses[i].Address.IP = ip

		if err := canonicalizeIP(&n.IPAM.Addresses[i].Address.IP); err != nil {
			return nil, "", fmt.Errorf("invalid address %d: %s", i, err)
		}

		if n.IPAM.Addresses[i].Address.IP.To4() != nil {
			n.IPAM.Addresses[i].Version = "4"
			numV4++
		} else {
			n.IPAM.Addresses[i].Version = "6"
			numV6++
		}
	}

	if envArgs != "" {
		e := IPAMEnvArgs{}
		err := types.LoadArgs(envArgs, &e)
		if err != nil {
			return nil, "", err
		}

		if e.IP != "" {
			for _, item := range strings.Split(string(e.IP), ",") {
				ipstr := strings.TrimSpace(item)

				ip, subnet, err := net.ParseCIDR(ipstr)
				if err != nil {
					return nil, "", fmt.Errorf("invalid CIDR %s: %s", ipstr, err)
				}

				addr := Address{Address: net.IPNet{IP: ip, Mask: subnet.Mask}}
				if addr.Address.IP.To4() != nil {
					addr.Version = "4"
					numV4++
				} else {
					addr.Version = "6"
					numV6++
				}
				n.IPAM.Addresses = append(n.IPAM.Addresses, addr)
			}
		}

		if e.GATEWAY != "" {
			for _, item := range strings.Split(string(e.GATEWAY), ",") {
				gwip := net.ParseIP(strings.TrimSpace(item))
				if gwip == nil {
					return nil, "", fmt.Errorf("invalid gateway address: %s", item)
				}

				for i := range n.IPAM.Addresses {
					if n.IPAM.Addresses[i].Address.Contains(gwip) {
						n.IPAM.Addresses[i].Gateway = gwip
					}
				}
			}
		}
	}

	// CNI spec 0.2.0 and below supported only one v4 and v6 address
	if numV4 > 1 || numV6 > 1 {
		for _, v := range types020.SupportedVersions {
			if n.CNIVersion == v {
				return nil, "", fmt.Errorf("CNI version %v does not support more than 1 address per family", n.CNIVersion)
			}
		}
	}

	// Copy net name into IPAM so not to drag Net struct around
	n.IPAM.Name = n.Name

	return n.IPAM, n.CNIVersion, nil
}

func cmdAdd(args *skel.CmdArgs) error {
	ipamConf, confVersion, err := LoadIPAMConfig(args.StdinData, args.Args)
	if err != nil {
		return err
	}

	result := &current.Result{}
	result.DNS = ipamConf.DNS
	result.Routes = ipamConf.Routes
	for _, v := range ipamConf.Addresses {
		result.IPs = append(result.IPs, &current.IPConfig{
			Version: v.Version,
			Address: v.Address,
			Gateway: v.Gateway})
	}

	return types.PrintResult(result, confVersion)
}

func cmdDel(args *skel.CmdArgs) error {
	// Nothing required because of no resource allocation in static plugin.
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"net"
	"os"
	"strings"
)

// Abort with an ENV-ERROR if the environment variable isn't set.
func getenvRequired(key string) string {
	value := os.Getenv(key)
	if value == "" {
		ptEnvError(fmt.Sprintf("no %s environment variable", key))
	}
	return value
}

// Escape a string so it contains no byte values over 127 and doesn't contain
// any of the characters '\x00', '\n', or '\\'.
func escape(s string) string {
	var buf bytes.Buffer
	for _, b := range []byte(s) {
		if b == '\n' {
			buf.WriteString("\\n")
		} else if b == '\\' {
			buf.WriteString("\\\\")
		} else if 0 < b && b < 128 {
			buf.WriteByte(b)
		} else {
			fmt.Fprintf(&buf, "\\x%02x", b)
		}
	}
	return buf.String()
}

func ptLine(keyword string, v ...string) {
	var buf bytes.Buffer
	buf.WriteString(keyword)
	for _, x := range v {
		buf.WriteString(" " + escape(x))
	}
	fmt.Println(buf.String())
}

func ptEnvError(msg string) {
	ptLine("ENV-ERROR", msg)
	os.Exit(1)
}

func ptVersionError(msg string) {
	ptLine("VERSION-ERROR", msg)
	os.Exit(1)
}

func ptCmethodError(methodName, msg string) {
	ptLine("CMETHOD-ERROR", methodName, msg)
	os.Exit(1)
}

func ptGetManagedTransportVer() string {
	const transportVersion = "1"
	for _, offered := range strings.Split(getenvRequired("TOR_PT_MANAGED_TRANSPORT_VER"), ",") {
		if offered == transportVersion {
			return offered
		}
	}
	return ""
}

func ptGetClientTransports(supported []string) []string {
	clientTransports := getenvRequired("TOR_PT_CLIENT_TRANSPORTS")
	if clientTransports == "*" {
		return supported
	}
	result := make([]string, 0)
	for _, requested := range strings.Split(clientTransports, ",") {
		for _, methodName := range supported {
			if requested == methodName {
				result = append(result, methodName)
				break
			}
		}
	}
	return result
}

func ptCmethod(name string, socks string, addr net.Addr) {
	ptLine("CMETHOD", name, socks, addr.String())
}

func ptCmethodsDone() {
	ptLine("CMETHODS", "DONE")
}

func ptClientSetup(methodNames []string) []string {
	ver := ptGetManagedTransportVer()
	if ver == "" {
		ptVersionError("no-version")
	} else {
		ptLine("VERSION", ver)
	}

	methods := ptGetClientTransports(methodNames)
	if len(methods) == 0 {
		ptCmethodsDone()
		os.Exit(1)
	}

	return methods
}
<commit_msg>Add server pluggable transports functions.<commit_after>package main

import (
	"bytes"
	"fmt"
	"net"
	"os"
	"strings"
)

// Abort with an ENV-ERROR if the environment variable isn't set.
func getenvRequired(key string) string {
	value := os.Getenv(key)
	if value == "" {
		ptEnvError(fmt.Sprintf("no %s environment variable", key))
	}
	return value
}

// Escape a string so it contains no byte values over 127 and doesn't contain
// any of the characters '\x00', '\n', or '\\'.
func escape(s string) string {
	var buf bytes.Buffer
	for _, b := range []byte(s) {
		if b == '\n' {
			buf.WriteString("\\n")
		} else if b == '\\' {
			buf.WriteString("\\\\")
		} else if 0 < b && b < 128 {
			buf.WriteByte(b)
		} else {
			fmt.Fprintf(&buf, "\\x%02x", b)
		}
	}
	return buf.String()
}

func ptLine(keyword string, v ...string) {
	var buf bytes.Buffer
	buf.WriteString(keyword)
	for _, x := range v {
		buf.WriteString(" " + escape(x))
	}
	fmt.Println(buf.String())
}

func ptEnvError(msg string) {
	ptLine("ENV-ERROR", msg)
	os.Exit(1)
}

func ptVersionError(msg string) {
	ptLine("VERSION-ERROR", msg)
	os.Exit(1)
}

func ptCmethodError(methodName, msg string) {
	ptLine("CMETHOD-ERROR", methodName, msg)
	os.Exit(1)
}

func ptGetManagedTransportVer() string {
	const transportVersion = "1"
	for _, offered := range strings.Split(getenvRequired("TOR_PT_MANAGED_TRANSPORT_VER"), ",") {
		if offered == transportVersion {
			return offered
		}
	}
	return ""
}

func ptGetClientTransports(supported []string) []string {
	clientTransports := getenvRequired("TOR_PT_CLIENT_TRANSPORTS")
	if clientTransports == "*" {
		return supported
	}
	result := make([]string, 0)
	for _, requested := range strings.Split(clientTransports, ",") {
		for _, methodName := range supported {
			if requested == methodName {
				result = append(result, methodName)
				break
			}
		}
	}
	return result
}

func ptCmethod(name string, socks string, addr net.Addr) {
	ptLine("CMETHOD", name, socks, addr.String())
}

func ptCmethodsDone() {
	ptLine("CMETHODS", "DONE")
}

func ptClientSetup(methodNames []string) []string {
	ver := ptGetManagedTransportVer()
	if ver == "" {
		ptVersionError("no-version")
	} else {
		ptLine("VERSION", ver)
	}

	methods := ptGetClientTransports(methodNames)
	if len(methods) == 0 {
		ptCmethodsDone()
		os.Exit(1)
	}

	return methods
}

func ptSmethodError(methodName, msg string) {
	ptLine("SMETHOD-ERROR", methodName, msg)
	os.Exit(1)
}

func ptSmethod(name string, addr net.Addr) {
	ptLine("SMETHOD", name, addr.String())
}

func ptSmethodsDone() {
	ptLine("SMETHODS", "DONE")
}

func resolveBindAddr(bindAddr string) (*net.TCPAddr, error) {
	addr, err := net.ResolveTCPAddr("tcp", bindAddr)
	if err == nil {
		return addr, nil
	}
	// Before the fixing of bug #7011, tor doesn't put brackets around IPv6
	// addresses. Split after the last colon, assuming it is a port
	// separator, and try adding the brackets.
	parts := strings.Split(bindAddr, ":")
	if len(parts) <= 2 {
		return nil, err
	}
	bindAddr = "[" + strings.Join(parts[:len(parts)-1], ":") + "]:" + parts[len(parts)-1]
	return net.ResolveTCPAddr("tcp", bindAddr)
}

type ptBindAddr struct {
	MethodName string
	Addr       *net.TCPAddr
}

func filterBindAddrs(addrs []ptBindAddr, supported []string) []ptBindAddr {
	var result []ptBindAddr

	for _, ba := range addrs {
		for _, methodName := range supported {
			if ba.MethodName == methodName {
				result = append(result, ba)
				break
			}
		}
	}

	return result
}

// Return a map from method names to bind addresses. The map is the contents of
// TOR_PT_SERVER_BINDADDR, with keys filtered by TOR_PT_SERVER_TRANSPORTS, and
// further filtered by methods that we know.
func ptGetServerBindAddrs(supported []string) []ptBindAddr {
	var result []ptBindAddr

	// Get the list of all requested bindaddrs.
	var serverBindAddr = getenvRequired("TOR_PT_SERVER_BINDADDR")
	for _, spec := range strings.Split(serverBindAddr, ",") {
		var bindAddr ptBindAddr

		parts := strings.SplitN(spec, "-", 2)
		if len(parts) != 2 {
			ptEnvError(fmt.Sprintf("TOR_PT_SERVER_BINDADDR: %q: doesn't contain \"-\"", spec))
		}
		bindAddr.MethodName = parts[0]
		addr, err := resolveBindAddr(parts[1])
		if err != nil {
			ptEnvError(fmt.Sprintf("TOR_PT_SERVER_BINDADDR: %q: %s", spec, err.Error()))
		}
		bindAddr.Addr = addr
		result = append(result, bindAddr)
	}

	// Filter by TOR_PT_SERVER_TRANSPORTS.
	serverTransports := getenvRequired("TOR_PT_SERVER_TRANSPORTS")
	if serverTransports != "*" {
		result = filterBindAddrs(result, strings.Split(serverTransports, ","))
	}

	// Finally filter by what we understand.
	result = filterBindAddrs(result, supported)

	return result
}

type ptServerInfo struct {
	BindAddrs []ptBindAddr
	OrAddr    *net.TCPAddr
}

func ptServerSetup(methodNames []string) ptServerInfo {
	var info ptServerInfo
	var err error

	ver := ptGetManagedTransportVer()
	if ver == "" {
		ptVersionError("no-version")
	} else {
		ptLine("VERSION", ver)
	}

	var orPort = getenvRequired("TOR_PT_ORPORT")
	info.OrAddr, err = net.ResolveTCPAddr("tcp", orPort)
	if err != nil {
		ptEnvError(fmt.Sprintf("cannot resolve TOR_PT_ORPORT %q: %s", orPort, err.Error()))
	}

	info.BindAddrs = ptGetServerBindAddrs(methodNames)
	if len(info.BindAddrs) == 0 {
		ptSmethodsDone()
		os.Exit(1)
	}

	return info
}
<|endoftext|>


<commit_before>// +build linux,!libdm_no_deferred_remove

package devicemapper

/*
#cgo LDFLAGS: -L. -ldevmapper
#include <libdevmapper.h>
*/
import "C"

// LibraryDeferredRemovalsupport is supported when statically linked.
const LibraryDeferredRemovalSupport = true

func dmTaskDeferredRemoveFct(task *CDmTask) int {
	return int(C.dm_task_deferred_remove((*C.struct_dm_task)(task)))
}

func dmTaskGetInfoWithDeferredFct(task *CDmTask, info *Info) int {
	Cinfo := C.struct_dm_info{}
	defer func() {
		info.Exists = int(Cinfo.exists)
		info.Suspended = int(Cinfo.suspended)
		info.LiveTable = int(Cinfo.live_table)
		info.InactiveTable = int(Cinfo.inactive_table)
		info.OpenCount = int32(Cinfo.open_count)
		info.EventNr = uint32(Cinfo.event_nr)
		info.Major = uint32(Cinfo.major)
		info.Minor = uint32(Cinfo.minor)
		info.ReadOnly = int(Cinfo.read_only)
		info.TargetCount = int32(Cinfo.target_count)
		info.DeferredRemove = int(Cinfo.deferred_remove)
	}()
	return int(C.dm_task_get_info((*C.struct_dm_task)(task), &Cinfo))
}
<commit_msg>pkg: devicemapper: fix typo in function declaration<commit_after>// +build linux,!libdm_no_deferred_remove

package devicemapper

/*
#cgo LDFLAGS: -L. -ldevmapper
#include <libdevmapper.h>
*/
import "C"

// LibraryDeferredRemovalsupport is supported when statically linked.
const LibraryDeferredRemovalSupport = true

func dmTaskDeferredRemoveFct(task *cdmTask) int {
	return int(C.dm_task_deferred_remove((*C.struct_dm_task)(task)))
}

func dmTaskGetInfoWithDeferredFct(task *cdmTask, info *Info) int {
	Cinfo := C.struct_dm_info{}
	defer func() {
		info.Exists = int(Cinfo.exists)
		info.Suspended = int(Cinfo.suspended)
		info.LiveTable = int(Cinfo.live_table)
		info.InactiveTable = int(Cinfo.inactive_table)
		info.OpenCount = int32(Cinfo.open_count)
		info.EventNr = uint32(Cinfo.event_nr)
		info.Major = uint32(Cinfo.major)
		info.Minor = uint32(Cinfo.minor)
		info.ReadOnly = int(Cinfo.read_only)
		info.TargetCount = int32(Cinfo.target_count)
		info.DeferredRemove = int(Cinfo.deferred_remove)
	}()
	return int(C.dm_task_get_info((*C.struct_dm_task)(task), &Cinfo))
}
<|endoftext|>


<commit_before>// Copyright 2017 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 main

import (
	"bytes"
	"fmt"
	"math/rand"
	"os"
	"runtime/debug"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/google/syzkaller/pkg/cover"
	"github.com/google/syzkaller/pkg/hash"
	"github.com/google/syzkaller/pkg/ipc"
	"github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/rpctype"
	"github.com/google/syzkaller/pkg/signal"
	"github.com/google/syzkaller/prog"
)

const (
	programLength = 30
)

// Proc represents a single fuzzing process (executor).
type Proc struct {
	fuzzer            *Fuzzer
	pid               int
	env               *ipc.Env
	rnd               *rand.Rand
	execOpts          *ipc.ExecOpts
	execOptsCover     *ipc.ExecOpts
	execOptsComps     *ipc.ExecOpts
	execOptsNoCollide *ipc.ExecOpts
}

func newProc(fuzzer *Fuzzer, pid int) (*Proc, error) {
	env, err := ipc.MakeEnv(fuzzer.config, pid)
	if err != nil {
		return nil, err
	}
	rnd := rand.New(rand.NewSource(time.Now().UnixNano() + int64(pid)*1e12))
	execOptsNoCollide := *fuzzer.execOpts
	execOptsNoCollide.Flags &= ^ipc.FlagCollide
	execOptsCover := execOptsNoCollide
	execOptsCover.Flags |= ipc.FlagCollectCover
	execOptsComps := execOptsNoCollide
	execOptsComps.Flags |= ipc.FlagCollectComps
	proc := &Proc{
		fuzzer:            fuzzer,
		pid:               pid,
		env:               env,
		rnd:               rnd,
		execOpts:          fuzzer.execOpts,
		execOptsCover:     &execOptsCover,
		execOptsComps:     &execOptsComps,
		execOptsNoCollide: &execOptsNoCollide,
	}
	return proc, nil
}

func (proc *Proc) loop() {
	generatePeriod := 100
	if proc.fuzzer.config.Flags&ipc.FlagSignal == 0 {
		// If we don't have real coverage signal, generate programs more frequently
		// because fallback signal is weak.
		generatePeriod = 10
	}
	for i := 0; ; i++ {
		item := proc.fuzzer.workQueue.dequeue()
		if item != nil {
			switch item := item.(type) {
			case *WorkTriage:
				proc.triageInput(item)
			case *WorkCandidate:
				proc.execute(proc.execOpts, item.p, item.flags, StatCandidate)
			case *WorkSmash:
				proc.smashInput(item)
			default:
				log.Fatalf("unknown work type: %#v", item)
			}
			continue
		}

		ct := proc.fuzzer.choiceTable
		corpus := proc.fuzzer.corpusSnapshot()
		if len(corpus) == 0 || i%generatePeriod == 0 {
			// Generate a new prog.
			p := proc.fuzzer.target.Generate(proc.rnd, programLength, ct)
			log.Logf(1, "#%v: generated", proc.pid)
			proc.execute(proc.execOpts, p, ProgNormal, StatGenerate)
		} else {
			// Mutate an existing prog.
			p := corpus[proc.rnd.Intn(len(corpus))].Clone()
			p.Mutate(proc.rnd, programLength, ct, corpus)
			log.Logf(1, "#%v: mutated", proc.pid)
			proc.execute(proc.execOpts, p, ProgNormal, StatFuzz)
		}
	}
}

func (proc *Proc) triageInput(item *WorkTriage) {
	log.Logf(1, "#%v: triaging type=%x", proc.pid, item.flags)

	call := item.p.Calls[item.call]
	inputSignal := signal.FromRaw(item.info.Signal, signalPrio(item.p.Target, call, &item.info))
	newSignal := proc.fuzzer.corpusSignalDiff(inputSignal)
	if newSignal.Empty() {
		return
	}
	log.Logf(3, "triaging input for %v (new signal=%v)", call.Meta.CallName, newSignal.Len())
	var inputCover cover.Cover
	const (
		signalRuns       = 3
		minimizeAttempts = 3
	)
	// Compute input coverage and non-flaky signal for minimization.
	notexecuted := 0
	for i := 0; i < signalRuns; i++ {
		info := proc.executeRaw(proc.execOptsCover, item.p, StatTriage)
		if len(info) == 0 || len(info[item.call].Signal) == 0 ||
			item.info.Errno == 0 && info[item.call].Errno != 0 {
			// The call was not executed or failed.
			notexecuted++
			if notexecuted > signalRuns/2+1 {
				return // if happens too often, give up
			}
			continue
		}
		inf := info[item.call]
		thisSignal := signal.FromRaw(inf.Signal, signalPrio(item.p.Target, call, &inf))
		newSignal = newSignal.Intersection(thisSignal)
		// Without !minimized check manager starts losing some considerable amount
		// of coverage after each restart. Mechanics of this are not completely clear.
		if newSignal.Empty() && item.flags&ProgMinimized == 0 {
			return
		}
		inputCover.Merge(inf.Cover)
	}
	if item.flags&ProgMinimized == 0 {
		item.p, item.call = prog.Minimize(item.p, item.call, false,
			func(p1 *prog.Prog, call1 int) bool {
				for i := 0; i < minimizeAttempts; i++ {
					info := proc.execute(proc.execOptsNoCollide, p1, ProgNormal, StatMinimize)
					if len(info) == 0 || len(info[call1].Signal) == 0 {
						continue // The call was not executed.
					}
					inf := info[call1]
					if item.info.Errno == 0 && inf.Errno != 0 {
						// Don't minimize calls from successful to unsuccessful.
						// Successful calls are much more valuable.
						return false
					}
					prio := signalPrio(p1.Target, p1.Calls[call1], &inf)
					thisSignal := signal.FromRaw(inf.Signal, prio)
					if newSignal.Intersection(thisSignal).Len() == newSignal.Len() {
						return true
					}
				}
				return false
			})
	}

	data := item.p.Serialize()
	sig := hash.Hash(data)

	log.Logf(2, "added new input for %v to corpus:\n%s", call.Meta.CallName, data)
	proc.fuzzer.sendInputToManager(rpctype.RPCInput{
		Call:   call.Meta.CallName,
		Prog:   data,
		Signal: inputSignal.Serialize(),
		Cover:  inputCover.Serialize(),
	})

	proc.fuzzer.addInputToCorpus(item.p, inputSignal, sig)

	if item.flags&ProgSmashed == 0 {
		proc.fuzzer.workQueue.enqueue(&WorkSmash{item.p, item.call})
	}
}

func (proc *Proc) smashInput(item *WorkSmash) {
	if proc.fuzzer.faultInjectionEnabled {
		proc.failCall(item.p, item.call)
	}
	if proc.fuzzer.comparisonTracingEnabled {
		proc.executeHintSeed(item.p, item.call)
	}
	corpus := proc.fuzzer.corpusSnapshot()
	for i := 0; i < 100; i++ {
		p := item.p.Clone()
		p.Mutate(proc.rnd, programLength, proc.fuzzer.choiceTable, corpus)
		log.Logf(1, "#%v: smash mutated", proc.pid)
		proc.execute(proc.execOpts, p, ProgNormal, StatSmash)
	}
}

func (proc *Proc) failCall(p *prog.Prog, call int) {
	for nth := 0; nth < 100; nth++ {
		log.Logf(1, "#%v: injecting fault into call %v/%v", proc.pid, call, nth)
		opts := *proc.execOpts
		opts.Flags |= ipc.FlagInjectFault
		opts.FaultCall = call
		opts.FaultNth = nth
		info := proc.executeRaw(&opts, p, StatSmash)
		if info != nil && len(info) > call && !info[call].FaultInjected {
			break
		}
	}
}

func (proc *Proc) executeHintSeed(p *prog.Prog, call int) {
	log.Logf(1, "#%v: collecting comparisons", proc.pid)
	// First execute the original program to dump comparisons from KCOV.
	info := proc.execute(proc.execOptsComps, p, ProgNormal, StatSeed)
	if info == nil {
		return
	}

	// Then mutate the initial program for every match between
	// a syscall argument and a comparison operand.
	// Execute each of such mutants to check if it gives new coverage.
	p.MutateWithHints(call, info[call].Comps, func(p *prog.Prog) {
		log.Logf(1, "#%v: executing comparison hint", proc.pid)
		proc.execute(proc.execOpts, p, ProgNormal, StatHint)
	})
}

func (proc *Proc) execute(execOpts *ipc.ExecOpts, p *prog.Prog, flags ProgTypes, stat Stat) []ipc.CallInfo {
	info := proc.executeRaw(execOpts, p, stat)
	for _, callIndex := range proc.fuzzer.checkNewSignal(p, info) {
		info := info[callIndex]
		// info.Signal points to the output shmem region, detach it before queueing.
		info.Signal = append([]uint32{}, info.Signal...)
		// None of the caller use Cover, so just nil it instead of detaching.
		// Note: triage input uses executeRaw to get coverage.
		info.Cover = nil
		proc.fuzzer.workQueue.enqueue(&WorkTriage{
			p:     p.Clone(),
			call:  callIndex,
			info:  info,
			flags: flags,
		})
	}
	return info
}

func (proc *Proc) executeRaw(opts *ipc.ExecOpts, p *prog.Prog, stat Stat) []ipc.CallInfo {
	if opts.Flags&ipc.FlagDedupCover == 0 {
		log.Fatalf("dedup cover is not enabled")
	}

	// Limit concurrency window and do leak checking once in a while.
	ticket := proc.fuzzer.gate.Enter()
	defer proc.fuzzer.gate.Leave(ticket)

	proc.logProgram(opts, p)
	try := 0
retry:
	atomic.AddUint64(&proc.fuzzer.stats[stat], 1)
	output, info, failed, hanged, err := proc.env.Exec(opts, p)
	if failed {
		// BUG in output should be recognized by manager.
		log.Logf(0, "BUG: executor-detected bug:\n%s", output)
		// Don't return any cover so that the input is not added to corpus.
		return nil
	}
	if err != nil {
		if _, ok := err.(ipc.ExecutorFailure); ok || try > 10 {
			log.Fatalf("executor failed %v times:\n%v", try, err)
		}
		try++
		log.Logf(4, "fuzzer detected executor failure='%v', retrying #%d\n", err, (try + 1))
		debug.FreeOSMemory()
		time.Sleep(time.Second)
		goto retry
	}
	log.Logf(2, "result failed=%v hanged=%v: %v\n", failed, hanged, string(output))
	return info
}

func (proc *Proc) logProgram(opts *ipc.ExecOpts, p *prog.Prog) {
	if proc.fuzzer.outputType == OutputNone {
		return
	}

	data := p.Serialize()
	strOpts := ""
	if opts.Flags&ipc.FlagInjectFault != 0 {
		strOpts = fmt.Sprintf(" (fault-call:%v fault-nth:%v)", opts.FaultCall, opts.FaultNth)
	}

	// The following output helps to understand what program crashed kernel.
	// It must not be intermixed.
	switch proc.fuzzer.outputType {
	case OutputStdout:
		now := time.Now()
		proc.fuzzer.logMu.Lock()
		fmt.Printf("%02v:%02v:%02v executing program %v%v:\n%s\n",
			now.Hour(), now.Minute(), now.Second(),
			proc.pid, strOpts, data)
		proc.fuzzer.logMu.Unlock()
	case OutputDmesg:
		fd, err := syscall.Open("/dev/kmsg", syscall.O_WRONLY, 0)
		if err == nil {
			buf := new(bytes.Buffer)
			fmt.Fprintf(buf, "syzkaller: executing program %v%v:\n%s\n",
				proc.pid, strOpts, data)
			syscall.Write(fd, buf.Bytes())
			syscall.Close(fd)
		}
	case OutputFile:
		f, err := os.Create(fmt.Sprintf("%v-%v.prog", proc.fuzzer.name, proc.pid))
		if err == nil {
			if strOpts != "" {
				fmt.Fprintf(f, "#%v\n", strOpts)
			}
			f.Write(data)
			f.Close()
		}
	default:
		log.Fatalf("unknown output type: %v", proc.fuzzer.outputType)
	}
}
<commit_msg>syz-fuzzer: generate programs more frequently if no coverage<commit_after>// Copyright 2017 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 main

import (
	"bytes"
	"fmt"
	"math/rand"
	"os"
	"runtime/debug"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/google/syzkaller/pkg/cover"
	"github.com/google/syzkaller/pkg/hash"
	"github.com/google/syzkaller/pkg/ipc"
	"github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/rpctype"
	"github.com/google/syzkaller/pkg/signal"
	"github.com/google/syzkaller/prog"
)

const (
	programLength = 30
)

// Proc represents a single fuzzing process (executor).
type Proc struct {
	fuzzer            *Fuzzer
	pid               int
	env               *ipc.Env
	rnd               *rand.Rand
	execOpts          *ipc.ExecOpts
	execOptsCover     *ipc.ExecOpts
	execOptsComps     *ipc.ExecOpts
	execOptsNoCollide *ipc.ExecOpts
}

func newProc(fuzzer *Fuzzer, pid int) (*Proc, error) {
	env, err := ipc.MakeEnv(fuzzer.config, pid)
	if err != nil {
		return nil, err
	}
	rnd := rand.New(rand.NewSource(time.Now().UnixNano() + int64(pid)*1e12))
	execOptsNoCollide := *fuzzer.execOpts
	execOptsNoCollide.Flags &= ^ipc.FlagCollide
	execOptsCover := execOptsNoCollide
	execOptsCover.Flags |= ipc.FlagCollectCover
	execOptsComps := execOptsNoCollide
	execOptsComps.Flags |= ipc.FlagCollectComps
	proc := &Proc{
		fuzzer:            fuzzer,
		pid:               pid,
		env:               env,
		rnd:               rnd,
		execOpts:          fuzzer.execOpts,
		execOptsCover:     &execOptsCover,
		execOptsComps:     &execOptsComps,
		execOptsNoCollide: &execOptsNoCollide,
	}
	return proc, nil
}

func (proc *Proc) loop() {
	generatePeriod := 100
	if proc.fuzzer.config.Flags&ipc.FlagSignal == 0 {
		// If we don't have real coverage signal, generate programs more frequently
		// because fallback signal is weak.
		generatePeriod = 2
	}
	for i := 0; ; i++ {
		item := proc.fuzzer.workQueue.dequeue()
		if item != nil {
			switch item := item.(type) {
			case *WorkTriage:
				proc.triageInput(item)
			case *WorkCandidate:
				proc.execute(proc.execOpts, item.p, item.flags, StatCandidate)
			case *WorkSmash:
				proc.smashInput(item)
			default:
				log.Fatalf("unknown work type: %#v", item)
			}
			continue
		}

		ct := proc.fuzzer.choiceTable
		corpus := proc.fuzzer.corpusSnapshot()
		if len(corpus) == 0 || i%generatePeriod == 0 {
			// Generate a new prog.
			p := proc.fuzzer.target.Generate(proc.rnd, programLength, ct)
			log.Logf(1, "#%v: generated", proc.pid)
			proc.execute(proc.execOpts, p, ProgNormal, StatGenerate)
		} else {
			// Mutate an existing prog.
			p := corpus[proc.rnd.Intn(len(corpus))].Clone()
			p.Mutate(proc.rnd, programLength, ct, corpus)
			log.Logf(1, "#%v: mutated", proc.pid)
			proc.execute(proc.execOpts, p, ProgNormal, StatFuzz)
		}
	}
}

func (proc *Proc) triageInput(item *WorkTriage) {
	log.Logf(1, "#%v: triaging type=%x", proc.pid, item.flags)

	call := item.p.Calls[item.call]
	inputSignal := signal.FromRaw(item.info.Signal, signalPrio(item.p.Target, call, &item.info))
	newSignal := proc.fuzzer.corpusSignalDiff(inputSignal)
	if newSignal.Empty() {
		return
	}
	log.Logf(3, "triaging input for %v (new signal=%v)", call.Meta.CallName, newSignal.Len())
	var inputCover cover.Cover
	const (
		signalRuns       = 3
		minimizeAttempts = 3
	)
	// Compute input coverage and non-flaky signal for minimization.
	notexecuted := 0
	for i := 0; i < signalRuns; i++ {
		info := proc.executeRaw(proc.execOptsCover, item.p, StatTriage)
		if len(info) == 0 || len(info[item.call].Signal) == 0 ||
			item.info.Errno == 0 && info[item.call].Errno != 0 {
			// The call was not executed or failed.
			notexecuted++
			if notexecuted > signalRuns/2+1 {
				return // if happens too often, give up
			}
			continue
		}
		inf := info[item.call]
		thisSignal := signal.FromRaw(inf.Signal, signalPrio(item.p.Target, call, &inf))
		newSignal = newSignal.Intersection(thisSignal)
		// Without !minimized check manager starts losing some considerable amount
		// of coverage after each restart. Mechanics of this are not completely clear.
		if newSignal.Empty() && item.flags&ProgMinimized == 0 {
			return
		}
		inputCover.Merge(inf.Cover)
	}
	if item.flags&ProgMinimized == 0 {
		item.p, item.call = prog.Minimize(item.p, item.call, false,
			func(p1 *prog.Prog, call1 int) bool {
				for i := 0; i < minimizeAttempts; i++ {
					info := proc.execute(proc.execOptsNoCollide, p1, ProgNormal, StatMinimize)
					if len(info) == 0 || len(info[call1].Signal) == 0 {
						continue // The call was not executed.
					}
					inf := info[call1]
					if item.info.Errno == 0 && inf.Errno != 0 {
						// Don't minimize calls from successful to unsuccessful.
						// Successful calls are much more valuable.
						return false
					}
					prio := signalPrio(p1.Target, p1.Calls[call1], &inf)
					thisSignal := signal.FromRaw(inf.Signal, prio)
					if newSignal.Intersection(thisSignal).Len() == newSignal.Len() {
						return true
					}
				}
				return false
			})
	}

	data := item.p.Serialize()
	sig := hash.Hash(data)

	log.Logf(2, "added new input for %v to corpus:\n%s", call.Meta.CallName, data)
	proc.fuzzer.sendInputToManager(rpctype.RPCInput{
		Call:   call.Meta.CallName,
		Prog:   data,
		Signal: inputSignal.Serialize(),
		Cover:  inputCover.Serialize(),
	})

	proc.fuzzer.addInputToCorpus(item.p, inputSignal, sig)

	if item.flags&ProgSmashed == 0 {
		proc.fuzzer.workQueue.enqueue(&WorkSmash{item.p, item.call})
	}
}

func (proc *Proc) smashInput(item *WorkSmash) {
	if proc.fuzzer.faultInjectionEnabled {
		proc.failCall(item.p, item.call)
	}
	if proc.fuzzer.comparisonTracingEnabled {
		proc.executeHintSeed(item.p, item.call)
	}
	corpus := proc.fuzzer.corpusSnapshot()
	for i := 0; i < 100; i++ {
		p := item.p.Clone()
		p.Mutate(proc.rnd, programLength, proc.fuzzer.choiceTable, corpus)
		log.Logf(1, "#%v: smash mutated", proc.pid)
		proc.execute(proc.execOpts, p, ProgNormal, StatSmash)
	}
}

func (proc *Proc) failCall(p *prog.Prog, call int) {
	for nth := 0; nth < 100; nth++ {
		log.Logf(1, "#%v: injecting fault into call %v/%v", proc.pid, call, nth)
		opts := *proc.execOpts
		opts.Flags |= ipc.FlagInjectFault
		opts.FaultCall = call
		opts.FaultNth = nth
		info := proc.executeRaw(&opts, p, StatSmash)
		if info != nil && len(info) > call && !info[call].FaultInjected {
			break
		}
	}
}

func (proc *Proc) executeHintSeed(p *prog.Prog, call int) {
	log.Logf(1, "#%v: collecting comparisons", proc.pid)
	// First execute the original program to dump comparisons from KCOV.
	info := proc.execute(proc.execOptsComps, p, ProgNormal, StatSeed)
	if info == nil {
		return
	}

	// Then mutate the initial program for every match between
	// a syscall argument and a comparison operand.
	// Execute each of such mutants to check if it gives new coverage.
	p.MutateWithHints(call, info[call].Comps, func(p *prog.Prog) {
		log.Logf(1, "#%v: executing comparison hint", proc.pid)
		proc.execute(proc.execOpts, p, ProgNormal, StatHint)
	})
}

func (proc *Proc) execute(execOpts *ipc.ExecOpts, p *prog.Prog, flags ProgTypes, stat Stat) []ipc.CallInfo {
	info := proc.executeRaw(execOpts, p, stat)
	for _, callIndex := range proc.fuzzer.checkNewSignal(p, info) {
		info := info[callIndex]
		// info.Signal points to the output shmem region, detach it before queueing.
		info.Signal = append([]uint32{}, info.Signal...)
		// None of the caller use Cover, so just nil it instead of detaching.
		// Note: triage input uses executeRaw to get coverage.
		info.Cover = nil
		proc.fuzzer.workQueue.enqueue(&WorkTriage{
			p:     p.Clone(),
			call:  callIndex,
			info:  info,
			flags: flags,
		})
	}
	return info
}

func (proc *Proc) executeRaw(opts *ipc.ExecOpts, p *prog.Prog, stat Stat) []ipc.CallInfo {
	if opts.Flags&ipc.FlagDedupCover == 0 {
		log.Fatalf("dedup cover is not enabled")
	}

	// Limit concurrency window and do leak checking once in a while.
	ticket := proc.fuzzer.gate.Enter()
	defer proc.fuzzer.gate.Leave(ticket)

	proc.logProgram(opts, p)
	try := 0
retry:
	atomic.AddUint64(&proc.fuzzer.stats[stat], 1)
	output, info, failed, hanged, err := proc.env.Exec(opts, p)
	if failed {
		// BUG in output should be recognized by manager.
		log.Logf(0, "BUG: executor-detected bug:\n%s", output)
		// Don't return any cover so that the input is not added to corpus.
		return nil
	}
	if err != nil {
		if _, ok := err.(ipc.ExecutorFailure); ok || try > 10 {
			log.Fatalf("executor failed %v times:\n%v", try, err)
		}
		try++
		log.Logf(4, "fuzzer detected executor failure='%v', retrying #%d\n", err, (try + 1))
		debug.FreeOSMemory()
		time.Sleep(time.Second)
		goto retry
	}
	log.Logf(2, "result failed=%v hanged=%v: %v\n", failed, hanged, string(output))
	return info
}

func (proc *Proc) logProgram(opts *ipc.ExecOpts, p *prog.Prog) {
	if proc.fuzzer.outputType == OutputNone {
		return
	}

	data := p.Serialize()
	strOpts := ""
	if opts.Flags&ipc.FlagInjectFault != 0 {
		strOpts = fmt.Sprintf(" (fault-call:%v fault-nth:%v)", opts.FaultCall, opts.FaultNth)
	}

	// The following output helps to understand what program crashed kernel.
	// It must not be intermixed.
	switch proc.fuzzer.outputType {
	case OutputStdout:
		now := time.Now()
		proc.fuzzer.logMu.Lock()
		fmt.Printf("%02v:%02v:%02v executing program %v%v:\n%s\n",
			now.Hour(), now.Minute(), now.Second(),
			proc.pid, strOpts, data)
		proc.fuzzer.logMu.Unlock()
	case OutputDmesg:
		fd, err := syscall.Open("/dev/kmsg", syscall.O_WRONLY, 0)
		if err == nil {
			buf := new(bytes.Buffer)
			fmt.Fprintf(buf, "syzkaller: executing program %v%v:\n%s\n",
				proc.pid, strOpts, data)
			syscall.Write(fd, buf.Bytes())
			syscall.Close(fd)
		}
	case OutputFile:
		f, err := os.Create(fmt.Sprintf("%v-%v.prog", proc.fuzzer.name, proc.pid))
		if err == nil {
			if strOpts != "" {
				fmt.Fprintf(f, "#%v\n", strOpts)
			}
			f.Write(data)
			f.Close()
		}
	default:
		log.Fatalf("unknown output type: %v", proc.fuzzer.outputType)
	}
}
<|endoftext|>


<commit_before>package gostardict

import (
	"compress/gzip"
	"io"
	"io/ioutil"
	"os"
	"strings"
)

// Dict implements in-memory dictionary
type Dict struct {
	buffer []byte
}

// GetSequence returns data at the given offset
func (d Dict) GetSequence(offset uint64, size uint64) []byte {
	return d.buffer[offset:(offset + size)]
}

// ReadDict reads dictionary into memory
func ReadDict(filename string, info *Info) (dict *Dict, err error) {
	reader, err := os.Open(filename)

	if err != nil {
		return
	}

	defer reader.Close()

	var r io.Reader

	if strings.HasSuffix(filename, ".dz") { // if file is compressed then read it from archive
		r, err = gzip.NewReader(reader)
	} else {
		r = reader
	}

	if err != nil {
		return
	}

	buffer, err := ioutil.ReadAll(r)

	dict = new(Dict)
	dict.buffer = buffer

	return
}
<commit_msg>error handling improvement<commit_after>package gostardict

import (
	"compress/gzip"
	"io"
	"io/ioutil"
	"os"
	"strings"
)

// Dict implements in-memory dictionary
type Dict struct {
	buffer []byte
}

// GetSequence returns data at the given offset
func (d Dict) GetSequence(offset uint64, size uint64) []byte {
	return d.buffer[offset:(offset + size)]
}

// ReadDict reads dictionary into memory
func ReadDict(filename string, info *Info) (dict *Dict, err error) {
	reader, err := os.Open(filename)

	if err != nil {
		return
	}

	defer reader.Close()

	var r io.Reader

	if strings.HasSuffix(filename, ".dz") { // if file is compressed then read it from archive
		r, err = gzip.NewReader(reader)
	} else {
		r = reader
	}

	if err != nil {
		return
	}

	buffer, err := ioutil.ReadAll(r)

	if err != nil {
		return
	}

	dict = new(Dict)
	dict.buffer = buffer

	return
}
<|endoftext|>


<commit_before>// Copyright (c) 2015, Álvaro Vilanova Vidal
// Use of this source code is governed by a BSD 3-Clause license that can be
// found in the LICENSE file.

package dropbox

import (
	"net/http"

	"testing"
)

func TestNewClient(t *testing.T) {
	c := NewClient(nil)

	if got, want := c.client, http.DefaultClient; got != want {
		t.Error("NewClient client is not http.DefaultClient")
	}
	if got, want := c.BaseURL.String(), defaultBaseURL; got != want {
		t.Errorf("NewClient BaseURL is %v, want %v", got, want)
	}
	if got, want := c.ContentURL.String(), defaultContentURL; got != want {
		t.Errorf("NewClient ContentURL is %v, want %v", got, want)
	}
	if got, want := c.UserAgent, userAgent; got != want {
		t.Errorf("NewClient UserAgent is %v, want %v", got, want)
	}
}

func TestCheckContentType(t *testing.T) {
	res := &http.Response{}
	if checkContentType(res, "") {
		t.Error("checkContentType on an empty response must return always false")
	}
	res.Header = http.Header{}
	res.Header.Add("Content-Type", "")
	if !checkContentType(res, "") {
		t.Error("checkContentType({Content-Type: \"\"}, \"\") returned false")
	}
	res.Header.Set("Content-Type", "  \n\t   application/json")
	if !checkContentType(res, "") {
		t.Error("checkContentType is no trimming strings")
	}
	res.Header.Set("Content-Type", "application/json;q=0.8")
	if !checkContentType(res, "application/json") {
		t.Error("checkContentType({Content-Type: \"application/json\"}, \"application/json\") returned false")
	}
}
<commit_msg>Added CheckResponse tests<commit_after>// Copyright (c) 2015, Álvaro Vilanova Vidal
// Use of this source code is governed by a BSD 3-Clause license that can be
// found in the LICENSE file.

package dropbox

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"testing/iotest"

	"testing"
)

func TestNewClient(t *testing.T) {
	c := NewClient(nil)

	if got, want := c.client, http.DefaultClient; got != want {
		t.Error("NewClient client is not http.DefaultClient")
	}
	if got, want := c.BaseURL.String(), defaultBaseURL; got != want {
		t.Errorf("NewClient BaseURL is %v, want %v", got, want)
	}
	if got, want := c.ContentURL.String(), defaultContentURL; got != want {
		t.Errorf("NewClient ContentURL is %v, want %v", got, want)
	}
	if got, want := c.UserAgent, userAgent; got != want {
		t.Errorf("NewClient UserAgent is %v, want %v", got, want)
	}
}

func TestCheckContentType(t *testing.T) {
	res := &http.Response{}
	if checkContentType(res, "") {
		t.Error("checkContentType on an empty response must return always false")
	}
	res.Header = http.Header{}
	res.Header.Add("Content-Type", "")
	if !checkContentType(res, "") {
		t.Error("checkContentType({Content-Type: \"\"}, \"\") returned false")
	}
	res.Header.Set("Content-Type", "  \n\t   application/json")
	if !checkContentType(res, "") {
		t.Error("checkContentType is no trimming strings")
	}
	res.Header.Set("Content-Type", "application/json;q=0.8")
	if !checkContentType(res, "application/json") {
		t.Error("checkContentType({Content-Type: \"application/json\"}, \"application/json\") returned false")
	}
}

func TestCheckResponse(t *testing.T) {
	res := &http.Response{}
	res.StatusCode = 200
	if got, want := checkResponse(res), error(nil); got != want {
		t.Errorf("checkResponse({StatusCode: 200}) returned %v, want %v", got, want)
	}
	res.StatusCode = 300
	if got, want := checkResponse(res), error(UnexpectedError{}); got != want {
		t.Errorf("checkResponse({StatusCode: 300}) returned %v, want %v", got, want)
	}
	res.StatusCode = 500
	if got, want := checkResponse(res), error(UnexpectedError{}); got != want {
		t.Errorf("checkResponse({StatusCode: 500}) returned %v, want %v", got, want)
	}
	res.Body = ioutil.NopCloser(bytes.NewBufferString("db-error"))
	if got, want := checkResponse(res), error(UnexpectedError{}); got != want {
		t.Errorf("checkResponse({StatusCode: 500, No Content-Type, Body:\"db-error\"}) returned %v, want %v", got, want)
	}
	res.Header = http.Header{}
	res.Header.Add("Content-Type", "text/plain")
	res.Body = ioutil.NopCloser(bytes.NewBufferString("db-error"))
	if got, want := checkResponse(res), error(Error{"db-error"}); got != want {
		t.Errorf("checkResponse({StatusCode: 500, \"text/plain\", Body:\"db-error\"}) returned %v, want %v", got, want)
	}
	res.Body = ioutil.NopCloser(iotest.TimeoutReader(bytes.NewBufferString("db-error")))
	if got, want := checkResponse(res), error(Error{"db-error"}); got == want {
		t.Errorf("checkResponse with an invalid text content must return an io error")
	}
	res.Header.Set("Content-Type", "application/json")
	res.Body = ioutil.NopCloser(bytes.NewBufferString("db-error"))
	if got, want := checkResponse(res), error(Error{"db-error"}); got == want {
		t.Errorf("checkResponse with an invalid JSON body must return an encoding error")
	}
	res.Body = ioutil.NopCloser(bytes.NewBufferString("{\"reason\":\"db-error\"}"))
	if got, want := checkResponse(res), error(Error{"db-error"}); got != want {
		t.Errorf("checkResponse({StatusCode: 500, \"application/json\", Body:\"{\"reason\":\"db-error\"}\"}) returned %v, want %v", got, want)
	}
}
<|endoftext|>


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

package gabi

import (
	"crypto/rand"
	"crypto/rsa"
	"encoding/xml"
	"math/big"
	"os"
)

const (
	//XMLHeader can be a used as the XML header when writing keys in XML format.
	XMLHeader = "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n"
	// MaxNumAttributes holds the (current) maximum number of attributes.
	MaxNumAttributes = 6
	// DefaultEpochLength is the default epoch length for public keys.
	DefaultEpochLength = 432000
)

// PrivateKey represents an issuer's private key.
type PrivateKey struct {
	XMLName xml.Name `xml:"http://www.zurich.ibm.com/security/idemix IssuerPrivateKey"`
	P       *big.Int `xml:"Elements>p"`
	Q       *big.Int `xml:"Elements>q"`
	PPrime  *big.Int `xml:"Elements>pPrime"`
	QPrime  *big.Int `xml:"Elements>qPrime"`
}

// NewPrivateKey creates a new issuer private key using the provided parameters.
func NewPrivateKey(p, q *big.Int) *PrivateKey {
	sk := PrivateKey{P: p, Q: q, PPrime: new(big.Int), QPrime: new(big.Int)}

	sk.PPrime.Sub(p, bigONE)
	sk.PPrime.Rsh(sk.PPrime, 1)

	sk.QPrime.Sub(q, bigONE)
	sk.QPrime.Rsh(sk.QPrime, 1)

	return &sk
}

// WriteToFile writes the private key to an xml file.
func (privk *PrivateKey) WriteToFile(filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	// Write the standard XML header
	_, err = f.Write([]byte(XMLHeader))
	if err != nil {
		return err
	}

	// And the actual xml body (with indentation)
	b, err := xml.MarshalIndent(privk, "", "   ")
	if err != nil {
		return err
	}
	_, err = f.Write(b)
	return err
}

// xmlBases is an auxiliary struct to encode/decode the odd way bases are
// represented in the xml representation of public keys
type xmlBases struct {
	Num   int      `xml:"num,attr"`
	Base0 *big.Int `xml:"Base_0"`
	Base1 *big.Int `xml:"Base_1"`
	Base2 *big.Int `xml:"Base_2"`
	Base3 *big.Int `xml:"Base_3"`
	Base4 *big.Int `xml:"Base_4"`
	Base5 *big.Int `xml:"Base_5"`
}

// xmlFeatures is an auxiliary struct to make the XML encoding/decoding a bit
// easier while keeping the struct for PublicKey somewhat simple.
type xmlFeatures struct {
	Epoch struct {
		Length int `xml:"length,attr"`
	}
}

// Bases is a type that is introduced to simplify the encoding/decoding of
// a PublicKey whilst using the xml support of Go's standard library.
type Bases []*big.Int

// UnmarshalXML is an internal function to simplify decoding a PublicKey from
// XML.
func (bl *Bases) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
	var t xmlBases

	if err := d.DecodeElement(&t, &start); err != nil {
		return err
	}
	*bl = Bases{t.Base0, t.Base1, t.Base2, t.Base3, t.Base4, t.Base5}
	return nil
}

// MarshalXML is an internal function to simplify encoding a PublicKey to XML.
func (bl *Bases) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	t := xmlBases{Num: len(*bl), Base0: (*bl)[0], Base1: (*bl)[1], Base2: (*bl)[2], Base3: (*bl)[3], Base4: (*bl)[4], Base5: (*bl)[5]}
	return e.EncodeElement(t, start)
}

// EpochLength is a type that is introduced to simplify the encoding/decoding of
// a PublicKey whilst using the xml support of Go's standard library.
type EpochLength int

// UnmarshalXML is an internal function to simplify decoding a PublicKey from
// XML.
func (el *EpochLength) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
	var t xmlFeatures

	if err := d.DecodeElement(&t, &start); err != nil {
		return err
	}
	*el = EpochLength(t.Epoch.Length)
	return nil
}

// MarshalXML is an internal function to simplify encoding a PublicKey to XML.
func (el *EpochLength) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	var t xmlFeatures
	t.Epoch.Length = int(*el)
	return e.EncodeElement(t, start)
}

// PublicKey represents an issuer's public key.
type PublicKey struct {
	XMLName     xml.Name          `xml:"http://www.zurich.ibm.com/security/idemix IssuerPublicKey"`
	N           *big.Int          `xml:"Elements>n"` // Modulus n
	Z           *big.Int          `xml:"Elements>Z"` // Generator Z
	S           *big.Int          `xml:"Elements>S"` // Generator S
	R           Bases             `xml:"Elements>Bases"`
	EpochLength EpochLength       `xml:"Features"`
	Params      *SystemParameters `xml:"-"`
}

// NewPublicKey creates and returns a new public key based on the provided parameters.
func NewPublicKey(N, Z, S *big.Int, R []*big.Int) *PublicKey {
	return &PublicKey{
		N:           N,
		Z:           Z,
		S:           S,
		R:           R,
		EpochLength: DefaultEpochLength,
		Params:      &DefaultSystemParameters,
	}
}

// WriteToFile writes the public key to an xml file.
func (pubk *PublicKey) WriteToFile(filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	// Write the standard XML header
	_, err = f.Write([]byte(XMLHeader))
	if err != nil {
		return err
	}

	// And the actual xml body (with indentation)
	b, err := xml.MarshalIndent(pubk, "", "   ")
	if err != nil {
		return err
	}
	_, err = f.Write(b)
	return err
}

// GenerateKeyPair generates a private/public keypair for an Issuer
func GenerateKeyPair(param *SystemParameters) (*PrivateKey, *PublicKey, error) {
	primeSize := param.Ln / 2
	// Let's just piggy-back on Go's own RSA key generation to generate our Issuer
	// keys.
	rsaPrivateKey, err := rsa.GenerateKey(rand.Reader, int(param.Ln))
	if err != nil {
		return nil, nil, err
	}

	priv := &PrivateKey{P: rsaPrivateKey.Primes[0], Q: rsaPrivateKey.Primes[1], PPrime: new(big.Int), QPrime: new(big.Int)}

	// compute p' and q'
	priv.PPrime.Sub(priv.P, bigONE)
	priv.PPrime.Rsh(priv.PPrime, 1)

	priv.QPrime.Sub(priv.Q, bigONE)
	priv.QPrime.Rsh(priv.QPrime, 1)

	// compute n
	pubk := &PublicKey{Params: param, EpochLength: DefaultEpochLength}
	pubk.N = new(big.Int).Mul(priv.P, priv.Q)

	// Find an acceptable value for S; we follow lead of the Silvia code here:
	// Pick a random l_n value and check whether it is a quadratic residue modulo n

	var s *big.Int
	for {
		s, err = randomBigInt(param.Ln)
		if err != nil {
			return nil, nil, err
		}
		// check if S \elem Z_n
		if s.Cmp(pubk.N) > 0 {
			continue
		}
		if legendreSymbol(s, priv.P) == 1 && legendreSymbol(s, priv.Q) == 1 {
			break
		}
	}

	pubk.S = s

	// Derive Z from S
	var x *big.Int
	for {
		x, _ = randomBigInt(primeSize)
		if x.Cmp(bigTWO) > 0 && x.Cmp(pubk.N) < 0 {
			break
		}
	}

	// Compute Z = S^x mod n
	pubk.Z = new(big.Int).Exp(pubk.S, x, pubk.N)

	// Derive R_i for i = 0...MaxNumAttributes from S
	pubk.R = make([]*big.Int, MaxNumAttributes)
	for i := 0; i < MaxNumAttributes; i++ {
		pubk.R[i] = new(big.Int)

		var x *big.Int
		for {
			x, _ = randomBigInt(primeSize)
			if x.Cmp(bigTWO) > 0 && x.Cmp(pubk.N) < 0 {
				break
			}
		}
		// Compute R_i = S^x mod n
		pubk.R[i].Exp(pubk.S, x, pubk.N)
	}

	return priv, pubk, nil
}
<commit_msg>Now actually use safe primes in key generation.<commit_after>// Copyright 2016 Maarten Everts. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gabi

import (
	"crypto/rand"
	"encoding/xml"
	"math/big"
	"os"
)

const (
	//XMLHeader can be a used as the XML header when writing keys in XML format.
	XMLHeader = "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n"
	// MaxNumAttributes holds the (current) maximum number of attributes.
	MaxNumAttributes = 6
	// DefaultEpochLength is the default epoch length for public keys.
	DefaultEpochLength = 432000
)

// PrivateKey represents an issuer's private key.
type PrivateKey struct {
	XMLName xml.Name `xml:"http://www.zurich.ibm.com/security/idemix IssuerPrivateKey"`
	P       *big.Int `xml:"Elements>p"`
	Q       *big.Int `xml:"Elements>q"`
	PPrime  *big.Int `xml:"Elements>pPrime"`
	QPrime  *big.Int `xml:"Elements>qPrime"`
}

// NewPrivateKey creates a new issuer private key using the provided parameters.
func NewPrivateKey(p, q *big.Int) *PrivateKey {
	sk := PrivateKey{P: p, Q: q, PPrime: new(big.Int), QPrime: new(big.Int)}

	sk.PPrime.Sub(p, bigONE)
	sk.PPrime.Rsh(sk.PPrime, 1)

	sk.QPrime.Sub(q, bigONE)
	sk.QPrime.Rsh(sk.QPrime, 1)

	return &sk
}

// WriteToFile writes the private key to an xml file.
func (privk *PrivateKey) WriteToFile(filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	// Write the standard XML header
	_, err = f.Write([]byte(XMLHeader))
	if err != nil {
		return err
	}

	// And the actual xml body (with indentation)
	b, err := xml.MarshalIndent(privk, "", "   ")
	if err != nil {
		return err
	}
	_, err = f.Write(b)
	return err
}

// xmlBases is an auxiliary struct to encode/decode the odd way bases are
// represented in the xml representation of public keys
type xmlBases struct {
	Num   int      `xml:"num,attr"`
	Base0 *big.Int `xml:"Base_0"`
	Base1 *big.Int `xml:"Base_1"`
	Base2 *big.Int `xml:"Base_2"`
	Base3 *big.Int `xml:"Base_3"`
	Base4 *big.Int `xml:"Base_4"`
	Base5 *big.Int `xml:"Base_5"`
}

// xmlFeatures is an auxiliary struct to make the XML encoding/decoding a bit
// easier while keeping the struct for PublicKey somewhat simple.
type xmlFeatures struct {
	Epoch struct {
		Length int `xml:"length,attr"`
	}
}

// Bases is a type that is introduced to simplify the encoding/decoding of
// a PublicKey whilst using the xml support of Go's standard library.
type Bases []*big.Int

// UnmarshalXML is an internal function to simplify decoding a PublicKey from
// XML.
func (bl *Bases) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
	var t xmlBases

	if err := d.DecodeElement(&t, &start); err != nil {
		return err
	}
	*bl = Bases{t.Base0, t.Base1, t.Base2, t.Base3, t.Base4, t.Base5}
	return nil
}

// MarshalXML is an internal function to simplify encoding a PublicKey to XML.
func (bl *Bases) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	t := xmlBases{Num: len(*bl), Base0: (*bl)[0], Base1: (*bl)[1], Base2: (*bl)[2], Base3: (*bl)[3], Base4: (*bl)[4], Base5: (*bl)[5]}
	return e.EncodeElement(t, start)
}

// EpochLength is a type that is introduced to simplify the encoding/decoding of
// a PublicKey whilst using the xml support of Go's standard library.
type EpochLength int

// UnmarshalXML is an internal function to simplify decoding a PublicKey from
// XML.
func (el *EpochLength) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
	var t xmlFeatures

	if err := d.DecodeElement(&t, &start); err != nil {
		return err
	}
	*el = EpochLength(t.Epoch.Length)
	return nil
}

// MarshalXML is an internal function to simplify encoding a PublicKey to XML.
func (el *EpochLength) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	var t xmlFeatures
	t.Epoch.Length = int(*el)
	return e.EncodeElement(t, start)
}

// PublicKey represents an issuer's public key.
type PublicKey struct {
	XMLName     xml.Name          `xml:"http://www.zurich.ibm.com/security/idemix IssuerPublicKey"`
	N           *big.Int          `xml:"Elements>n"` // Modulus n
	Z           *big.Int          `xml:"Elements>Z"` // Generator Z
	S           *big.Int          `xml:"Elements>S"` // Generator S
	R           Bases             `xml:"Elements>Bases"`
	EpochLength EpochLength       `xml:"Features"`
	Params      *SystemParameters `xml:"-"`
}

// NewPublicKey creates and returns a new public key based on the provided parameters.
func NewPublicKey(N, Z, S *big.Int, R []*big.Int) *PublicKey {
	return &PublicKey{
		N:           N,
		Z:           Z,
		S:           S,
		R:           R,
		EpochLength: DefaultEpochLength,
		Params:      &DefaultSystemParameters,
	}
}

// WriteToFile writes the public key to an xml file.
func (pubk *PublicKey) WriteToFile(filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	// Write the standard XML header
	_, err = f.Write([]byte(XMLHeader))
	if err != nil {
		return err
	}

	// And the actual xml body (with indentation)
	b, err := xml.MarshalIndent(pubk, "", "   ")
	if err != nil {
		return err
	}
	_, err = f.Write(b)
	return err
}

// randomSafePrime produces a safe prime of the requested number of bits
func randomSafePrime(bits int) (*big.Int, error) {
	p2 := new(big.Int)
	for {
		p, err := rand.Prime(rand.Reader, bits)
		if err != nil {
			return nil, err
		}
		p2.Rsh(p, 1) // p2 = (p - 1)/2
		if p2.ProbablyPrime(20) {
			return p, nil
		}
	}
}

// GenerateKeyPair generates a private/public keypair for an Issuer
func GenerateKeyPair(param *SystemParameters) (*PrivateKey, *PublicKey, error) {
	primeSize := param.Ln / 2

	// p and q need to be safe primes
	p, err := randomSafePrime(int(primeSize))
	if err != nil {
		return nil, nil, err
	}

	q, err := randomSafePrime(int(primeSize))
	if err != nil {
		return nil, nil, err
	}

	priv := &PrivateKey{P: p, Q: q, PPrime: new(big.Int), QPrime: new(big.Int)}

	// compute p' and q'
	priv.PPrime.Sub(priv.P, bigONE)
	priv.PPrime.Rsh(priv.PPrime, 1)

	priv.QPrime.Sub(priv.Q, bigONE)
	priv.QPrime.Rsh(priv.QPrime, 1)

	// compute n
	pubk := &PublicKey{Params: param, EpochLength: DefaultEpochLength}
	pubk.N = new(big.Int).Mul(priv.P, priv.Q)

	// Find an acceptable value for S; we follow lead of the Silvia code here:
	// Pick a random l_n value and check whether it is a quadratic residue modulo n

	var s *big.Int
	for {
		s, err = randomBigInt(param.Ln)
		if err != nil {
			return nil, nil, err
		}
		// check if S \elem Z_n
		if s.Cmp(pubk.N) > 0 {
			continue
		}
		if legendreSymbol(s, priv.P) == 1 && legendreSymbol(s, priv.Q) == 1 {
			break
		}
	}

	pubk.S = s

	// Derive Z from S
	var x *big.Int
	for {
		x, _ = randomBigInt(primeSize)
		if x.Cmp(bigTWO) > 0 && x.Cmp(pubk.N) < 0 {
			break
		}
	}

	// Compute Z = S^x mod n
	pubk.Z = new(big.Int).Exp(pubk.S, x, pubk.N)

	// Derive R_i for i = 0...MaxNumAttributes from S
	pubk.R = make([]*big.Int, MaxNumAttributes)
	for i := 0; i < MaxNumAttributes; i++ {
		pubk.R[i] = new(big.Int)

		var x *big.Int
		for {
			x, _ = randomBigInt(primeSize)
			if x.Cmp(bigTWO) > 0 && x.Cmp(pubk.N) < 0 {
				break
			}
		}
		// Compute R_i = S^x mod n
		pubk.R[i].Exp(pubk.S, x, pubk.N)
	}

	return priv, pubk, nil
}
<|endoftext|>


<commit_before>package formatting

import (
	"io"
	"regexp"

	"github.com/AlexanderThaller/lablog/src/data"
)

func Notes(writer io.Writer, indent int, notes []data.Note) {
	for _, note := range notes {
		io.WriteString(writer, HeaderIndent(indent)+" ")
		io.WriteString(writer, note.TimeStamp.String()+"\n")

		NotesValue(writer, note.Value, indent+1)
		io.WriteString(writer, "\n")
	}
}

func NotesValue(writer io.Writer, value string, indent int) {
	indentchar := HeaderIndent(indent)
	indentreg, _ := regexp.Compile("(?m)^=")
	value = indentreg.ReplaceAllString(value, indentchar)

	io.WriteString(writer, value)
	io.WriteString(writer, "\n")
}

func ProjectNotes(writer io.Writer, indent int, project *data.Project) {
	if len(project.Notes()) == 0 {
		return
	}

	HeaderProject(writer, indent+1, project)
	Notes(writer, indent+2, project.Notes())
}
<commit_msg>Will not print empty notes anymore.<commit_after>package formatting

import (
	"io"
	"regexp"

	"github.com/AlexanderThaller/lablog/src/data"
)

func Notes(writer io.Writer, indent int, notes []data.Note) {
	for _, note := range notes {
		if note.Value == "" {
			continue
		}

		io.WriteString(writer, HeaderIndent(indent)+" ")
		io.WriteString(writer, note.TimeStamp.String()+"\n")

		NotesValue(writer, note.Value, indent+1)
		io.WriteString(writer, "\n")
	}
}

func NotesValue(writer io.Writer, value string, indent int) {
	indentchar := HeaderIndent(indent)
	indentreg, _ := regexp.Compile("(?m)^=")
	value = indentreg.ReplaceAllString(value, indentchar)

	io.WriteString(writer, value)
	io.WriteString(writer, "\n")
}

func ProjectNotes(writer io.Writer, indent int, project *data.Project) {
	if len(project.Notes()) == 0 {
		return
	}

	HeaderProject(writer, indent+1, project)
	Notes(writer, indent+2, project.Notes())
}
<|endoftext|>


<commit_before>package pqtgogen

import (
	"go/format"
	"io"

	"github.com/piotrkowalczuk/pqt"
	"github.com/piotrkowalczuk/pqt/internal/gogen"
	"github.com/piotrkowalczuk/pqt/internal/print"
)

// Component bits represents single component that can be generated by the generator.
type Component uint64

const (
	// ComponentInsert represents Insert method of a repository.
	ComponentInsert Component = 1 << (64 - 1 - iota)
	// ComponentFind represents Find method of a repository.
	ComponentFind
	// ComponentUpdate represents Update method of a repository.
	ComponentUpdate
	// ComponentUpsert represents Upsert method of a repository.
	ComponentUpsert
	// ComponentCount represents Count method of a repository.
	ComponentCount
	// ComponentDelete represents Delete method of a repository.
	ComponentDelete
	// ComponentHelpers represents all helpers.
	ComponentHelpers

	// ComponentRepository is a bit mask that group all repository methods.
	ComponentRepository = ComponentInsert | ComponentFind | ComponentUpdate | ComponentUpsert | ComponentCount | ComponentDelete
	// ComponentAll is a bit mask that groups all components.
	ComponentAll = ComponentRepository | ComponentHelpers
)

// Generator ...
type Generator struct {
	// Version represents
	Version    float64
	Pkg        string
	Imports    []string
	Plugins    []Plugin
	Components Component

	g *gogen.Generator
	p *print.Printer
}

// Generate generates formatted Go code for given schema.
func (g *Generator) Generate(s *pqt.Schema) ([]byte, error) {
	if err := g.generate(s); err != nil {
		return nil, err
	}

	return format.Source(g.p.Bytes())
}

// GenerateTo works like generate but it writes into io Writer instead.
func (g *Generator) GenerateTo(s *pqt.Schema, w io.Writer) error {
	if err := g.generate(s); err != nil {
		return err
	}

	buf, err := format.Source(g.p.Bytes())
	if err != nil {
		return err
	}
	_, err = w.Write(buf)
	return err
}

func (g *Generator) generate(s *pqt.Schema) error {
	g.g = &gogen.Generator{
		Version: g.Version,
	}
	for _, p := range g.Plugins {
		g.g.Plugins = append(g.g.Plugins, p)
	}
	g.p = &g.g.Printer

	g.g.Package(g.Pkg)
	g.g.Imports(s, "github.com/m4rw3r/uuid")
	if g.Components&ComponentRepository != 0 {
		g.g.Funcs()
		g.g.NewLine()
	}
	if g.Components&ComponentFind != 0 || g.Components&ComponentCount != 0 || g.Components&ComponentHelpers != 0 {
		g.g.Interfaces()
		g.g.NewLine()
	}
	if g.Components&ComponentFind != 0 || g.Components&ComponentCount != 0 {
		g.g.JoinClause()
		g.g.NewLine()
	}
	for _, t := range s.Tables {
		g.g.Constraints(t)
		g.g.NewLine()
		g.g.Columns(t)
		g.g.NewLine()
		g.g.Entity(t)
		g.g.NewLine()
		g.g.EntityProp(t)
		g.g.NewLine()
		g.g.EntityProps(t)
		g.g.NewLine()
		if g.Components&ComponentHelpers != 0 {
			g.g.ScanRows(t)
			g.g.NewLine()
		}
		if g.Components&ComponentFind != 0 || g.Components&ComponentCount != 0 {
			g.g.Iterator(t)
			g.g.NewLine()
			g.g.Criteria(t)
			g.g.NewLine()
			g.g.Operand(t)
			g.g.NewLine()
			g.g.FindExpr(t)
			g.g.NewLine()
			g.g.Join(t)
			g.g.NewLine()
		}
		if g.Components&ComponentCount != 0 {
			g.g.CountExpr(t)
			g.g.NewLine()
		}
		if g.Components&ComponentUpdate != 0 || g.Components&ComponentUpsert != 0 {
			g.g.Patch(t)
			g.g.NewLine()
		}
		if g.Components&ComponentRepository != 0 {
			g.g.Repository(t)
			g.g.NewLine()

			if g.Components&ComponentInsert != 0 {
				g.g.RepositoryInsertQuery(t)
				g.g.NewLine()
				g.g.RepositoryInsert(t)
				g.g.NewLine()
			}
			if g.Components&ComponentFind != 0 {
				g.g.WhereClause(t)
				g.g.NewLine()
				g.g.RepositoryFindQuery(t)
				g.g.NewLine()
				g.g.RepositoryFind(t)
				g.g.NewLine()
				g.g.RepositoryFindIter(t)
				g.g.NewLine()
				g.g.RepositoryFindOneByPrimaryKey(t)
				g.g.NewLine()
				g.g.RepositoryFindOneByUniqueConstraint(t)
				g.g.NewLine()
			}
			if g.Components&ComponentUpdate != 0 {
				g.g.RepositoryUpdateOneByPrimaryKeyQuery(t)
				g.g.NewLine()
				g.g.RepositoryUpdateOneByPrimaryKey(t)
				g.g.NewLine()
				g.g.RepositoryUpdateOneByUniqueConstraintQuery(t)
				g.g.NewLine()
				g.g.RepositoryUpdateOneByUniqueConstraint(t)
				g.g.NewLine()
			}
			if g.Components&ComponentUpsert != 0 {
				g.g.RepositoryUpsertQuery(t)
				g.g.NewLine()
				g.g.RepositoryUpsert(t)
				g.g.NewLine()
			}
			if g.Components&ComponentCount != 0 {
				g.g.RepositoryCount(t)
				g.g.NewLine()
			}
			if g.Components&ComponentDelete != 0 {
				g.g.RepositoryDeleteOneByPrimaryKey(t)
				g.g.NewLine()
			}
		}
	}
	g.g.Statics(s)
	g.g.NewLine()

	return g.p.Err
}
<commit_msg>pqtgo - more comments #2<commit_after>package pqtgogen

import (
	"go/format"
	"io"

	"github.com/piotrkowalczuk/pqt"
	"github.com/piotrkowalczuk/pqt/internal/gogen"
	"github.com/piotrkowalczuk/pqt/internal/print"
)

// Component bits represents single component that can be generated by the generator.
type Component uint64

const (
	// ComponentInsert represents Insert method of a repository.
	ComponentInsert Component = 1 << (64 - 1 - iota)
	// ComponentFind represents Find method of a repository.
	ComponentFind
	// ComponentUpdate represents Update method of a repository.
	ComponentUpdate
	// ComponentUpsert represents Upsert method of a repository.
	ComponentUpsert
	// ComponentCount represents Count method of a repository.
	ComponentCount
	// ComponentDelete represents Delete method of a repository.
	ComponentDelete
	// ComponentHelpers represents all helpers.
	ComponentHelpers

	// ComponentRepository is a bit mask that group all repository methods.
	ComponentRepository = ComponentInsert | ComponentFind | ComponentUpdate | ComponentUpsert | ComponentCount | ComponentDelete
	// ComponentAll is a bit mask that groups all components.
	ComponentAll = ComponentRepository | ComponentHelpers
)

// Generator ...
type Generator struct {
	// Version represents Postgres database version code will run against.
	Version float64
	// Pkg is the name of package code is generated into.
	// By default it's "main".
	Pkg string
	// Imports allow to pass additional import paths that will be added into generated code.
	Imports []string
	// Plugins that generator will use during generation.
	Plugins []Plugin
	// Components ...
	Components Component

	g *gogen.Generator
	p *print.Printer
}

// Generate generates formatted Go code for given schema.
func (g *Generator) Generate(s *pqt.Schema) ([]byte, error) {
	if err := g.generate(s); err != nil {
		return nil, err
	}

	return format.Source(g.p.Bytes())
}

// GenerateTo works like generate but it writes into io Writer instead.
func (g *Generator) GenerateTo(s *pqt.Schema, w io.Writer) error {
	if err := g.generate(s); err != nil {
		return err
	}

	buf, err := format.Source(g.p.Bytes())
	if err != nil {
		return err
	}
	_, err = w.Write(buf)
	return err
}

func (g *Generator) generate(s *pqt.Schema) error {
	g.g = &gogen.Generator{
		Version: g.Version,
	}
	for _, p := range g.Plugins {
		g.g.Plugins = append(g.g.Plugins, p)
	}
	g.p = &g.g.Printer

	g.g.Package(g.Pkg)
	g.g.Imports(s, "github.com/m4rw3r/uuid")
	if g.Components&ComponentRepository != 0 {
		g.g.Funcs()
		g.g.NewLine()
	}
	if g.Components&ComponentFind != 0 || g.Components&ComponentCount != 0 || g.Components&ComponentHelpers != 0 {
		g.g.Interfaces()
		g.g.NewLine()
	}
	if g.Components&ComponentFind != 0 || g.Components&ComponentCount != 0 {
		g.g.JoinClause()
		g.g.NewLine()
	}
	for _, t := range s.Tables {
		g.g.Constraints(t)
		g.g.NewLine()
		g.g.Columns(t)
		g.g.NewLine()
		g.g.Entity(t)
		g.g.NewLine()
		g.g.EntityProp(t)
		g.g.NewLine()
		g.g.EntityProps(t)
		g.g.NewLine()
		if g.Components&ComponentHelpers != 0 {
			g.g.ScanRows(t)
			g.g.NewLine()
		}
		if g.Components&ComponentFind != 0 || g.Components&ComponentCount != 0 {
			g.g.Iterator(t)
			g.g.NewLine()
			g.g.Criteria(t)
			g.g.NewLine()
			g.g.Operand(t)
			g.g.NewLine()
			g.g.FindExpr(t)
			g.g.NewLine()
			g.g.Join(t)
			g.g.NewLine()
		}
		if g.Components&ComponentCount != 0 {
			g.g.CountExpr(t)
			g.g.NewLine()
		}
		if g.Components&ComponentUpdate != 0 || g.Components&ComponentUpsert != 0 {
			g.g.Patch(t)
			g.g.NewLine()
		}
		if g.Components&ComponentRepository != 0 {
			g.g.Repository(t)
			g.g.NewLine()

			if g.Components&ComponentInsert != 0 {
				g.g.RepositoryInsertQuery(t)
				g.g.NewLine()
				g.g.RepositoryInsert(t)
				g.g.NewLine()
			}
			if g.Components&ComponentFind != 0 {
				g.g.WhereClause(t)
				g.g.NewLine()
				g.g.RepositoryFindQuery(t)
				g.g.NewLine()
				g.g.RepositoryFind(t)
				g.g.NewLine()
				g.g.RepositoryFindIter(t)
				g.g.NewLine()
				g.g.RepositoryFindOneByPrimaryKey(t)
				g.g.NewLine()
				g.g.RepositoryFindOneByUniqueConstraint(t)
				g.g.NewLine()
			}
			if g.Components&ComponentUpdate != 0 {
				g.g.RepositoryUpdateOneByPrimaryKeyQuery(t)
				g.g.NewLine()
				g.g.RepositoryUpdateOneByPrimaryKey(t)
				g.g.NewLine()
				g.g.RepositoryUpdateOneByUniqueConstraintQuery(t)
				g.g.NewLine()
				g.g.RepositoryUpdateOneByUniqueConstraint(t)
				g.g.NewLine()
			}
			if g.Components&ComponentUpsert != 0 {
				g.g.RepositoryUpsertQuery(t)
				g.g.NewLine()
				g.g.RepositoryUpsert(t)
				g.g.NewLine()
			}
			if g.Components&ComponentCount != 0 {
				g.g.RepositoryCount(t)
				g.g.NewLine()
			}
			if g.Components&ComponentDelete != 0 {
				g.g.RepositoryDeleteOneByPrimaryKey(t)
				g.g.NewLine()
			}
		}
	}
	g.g.Statics(s)
	g.g.NewLine()

	return g.p.Err
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"text/template"
	"time"
)

// DRAC contains all of the information required
// to connect to a Dell DRAC KVM
type DRAC struct {
	Host     string
	Username string
	Password string
	Version  int
}

// Templates is a map of each viewer.jnlp template for
// the various Dell iDRAC versions, keyed by version number
var Templates = map[int]string{
	6: viewer6,
	7: viewer7,
}

// GetVersion attempts to detect the iDRAC version by checking
// if various known libraries are available via HTTP GET requests.
// Retursn the version if found, or -1 if unknown
func (d *DRAC) GetVersion() int {

	log.Print("Detecting iDRAC version...")

	version := -1

	transport := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
		Dial: func(netw, addr string) (net.Conn, error) {
			deadline := time.Now().Add(5 * time.Second)
			c, err := net.DialTimeout(netw, addr, time.Second*5)
			if err != nil {
				return nil, err
			}
			c.SetDeadline(deadline)
			return c, nil
		},
	}

	client := &http.Client{
		Transport: transport,
	}

	// Check for iDRAC7 specific libs
	if response, err := client.Get("https://" + d.Host + "/software/avctKVMIOMac64.jar"); err == nil {
		response.Body.Close()
		if response.StatusCode == 200 {
			return 7
		}
	}

	// Check for iDRAC6 specific libs
	if response, err := client.Get("https://" + d.Host + "/software/jpcsc.jar"); err == nil {
		response.Body.Close()
		if response.StatusCode == 200 {
			return 6
		}
	}

	return version

}

// Viewer returns a viewer.jnlp template filled out with the
// necessary details to connect to a particular DRAC host
func (d *DRAC) Viewer() (string, error) {

	var version int

	// Check we have a valid DRAC viewer template for this DRAC version
	if d.Version < 0 {
		version = d.GetVersion()
	} else {
		version = d.Version
	}
	if version < 0 {
		return "", errors.New("unable to detect DRAC version")
	}

	log.Printf("Found iDRAC version %d", version)

	if _, ok := Templates[version]; !ok {
		msg := fmt.Sprintf("no support for DRAC v%d", version)
		return "", errors.New(msg)
	}

	// Generate a JNLP viewer from the template
	// Injecting the host/user/pass information
	buff := bytes.NewBufferString("")
	err := template.Must(template.New("viewer").Parse(Templates[version])).Execute(buff, d)
	return buff.String(), err

}
<commit_msg>Use 'defer' to call response.Body.Close()<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"text/template"
	"time"
)

// DRAC contains all of the information required
// to connect to a Dell DRAC KVM
type DRAC struct {
	Host     string
	Username string
	Password string
	Version  int
}

// Templates is a map of each viewer.jnlp template for
// the various Dell iDRAC versions, keyed by version number
var Templates = map[int]string{
	6: viewer6,
	7: viewer7,
}

// GetVersion attempts to detect the iDRAC version by checking
// if various known libraries are available via HTTP GET requests.
// Retursn the version if found, or -1 if unknown
func (d *DRAC) GetVersion() int {

	log.Print("Detecting iDRAC version...")

	version := -1

	transport := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
		Dial: func(netw, addr string) (net.Conn, error) {
			deadline := time.Now().Add(5 * time.Second)
			c, err := net.DialTimeout(netw, addr, time.Second*5)
			if err != nil {
				return nil, err
			}
			c.SetDeadline(deadline)
			return c, nil
		},
	}

	client := &http.Client{
		Transport: transport,
	}

	// Check for iDRAC7 specific libs
	if response, err := client.Get("https://" + d.Host + "/software/avctKVMIOMac64.jar"); err == nil {
		defer response.Body.Close()
		if response.StatusCode == 200 {
			return 7
		}
	}

	// Check for iDRAC6 specific libs
	if response, err := client.Get("https://" + d.Host + "/software/jpcsc.jar"); err == nil {
		defer response.Body.Close()
		if response.StatusCode == 200 {
			return 6
		}
	}

	return version

}

// Viewer returns a viewer.jnlp template filled out with the
// necessary details to connect to a particular DRAC host
func (d *DRAC) Viewer() (string, error) {

	var version int

	// Check we have a valid DRAC viewer template for this DRAC version
	if d.Version < 0 {
		version = d.GetVersion()
	} else {
		version = d.Version
	}
	if version < 0 {
		return "", errors.New("unable to detect DRAC version")
	}

	log.Printf("Found iDRAC version %d", version)

	if _, ok := Templates[version]; !ok {
		msg := fmt.Sprintf("no support for DRAC v%d", version)
		return "", errors.New(msg)
	}

	// Generate a JNLP viewer from the template
	// Injecting the host/user/pass information
	buff := bytes.NewBufferString("")
	err := template.Must(template.New("viewer").Parse(Templates[version])).Execute(buff, d)
	return buff.String(), err

}
<|endoftext|>


<commit_before>package ioprogress

import (
	"fmt"
	"io"
	"os"
	"strings"
)

// DrawFunc is the callback type for drawing progress.
type DrawFunc func(int64, int64) error

// DrawTextFormatFunc is a callback used by DrawFuncs that draw text in
// order to format the text into some more human friendly format.
type DrawTextFormatFunc func(int64, int64) string

var defaultDrawFunc DrawFunc

func init() {
	defaultDrawFunc = DrawTerminal(os.Stdout)
}

// DrawTerminal returns a DrawFunc that draws a progress bar to an io.Writer
// that is assumed to be a terminal (and therefore respects carriage returns).
func DrawTerminal(w io.Writer) DrawFunc {
	return DrawTerminalf(w, func(progress, total int64) string {
		return fmt.Sprintf("%d/%d", progress, total)
	})
}

// DrawTerminalf returns a DrawFunc that draws a progress bar to an io.Writer
// that is formatted with the given formatting function.
func DrawTerminalf(w io.Writer, f DrawTextFormatFunc) DrawFunc {
	var maxLength int

	return func(progress, total int64) error {
		if progress == -1 && total == -1 {
			_, err := fmt.Fprintf(w, "\n")
			return err
		}

		// Make sure we pad it to the max length we've ever drawn so that
		// we don't have trailing characters.
		line := f(progress, total)
		if len(line) < maxLength {
			line = fmt.Sprintf(
				"%s%s",
				line,
				strings.Repeat(" ", maxLength-len(line)))
		}
		maxLength = len(line)

		_, err := fmt.Fprint(w, line+"\r")
		return err
	}
}

var byteUnits = []string{"B", "KB", "MB", "GB", "TB", "PB"}

// DrawTextFormatBytes is a DrawTextFormatFunc that formats the progress
// and total into human-friendly byte formats.
func DrawTextFormatBytes(progress, total int64) string {
	return fmt.Sprintf("%s/%s", ByteUnitStr(progress), ByteUnitStr(total))
}

// DrawTextFormatBar returns a DrawTextFormatFunc that draws a progress
// bar with the given width (in characters). This can be used in conjunction
// with another DrawTextFormatFunc to create a progress bar with bytes, for
// example:
//
//     bar := DrawTextFormatBar(20)
//     func(progress, total int64) string {
//         return fmt.Sprintf(
//           "%s %s",
//           bar(progress, total),
//           DrawTextFormatBytes(progress, total))
//     }
//
func DrawTextFormatBar(width int64) DrawTextFormatFunc {
	width -= 2

	return func(progress, total int64) string {
		current := int64((float64(progress) / float64(total)) * float64(width))
		return fmt.Sprintf(
			"[%s%s]",
			strings.Repeat("=", int(current)),
			strings.Repeat(" ", int(width-current)))
	}
}

// ByteUnitStr pretty prints a number of bytes.
func ByteUnitStr(n int64) string {
	var unit string
	size := float64(n)
	for i := 1; i < len(byteUnits); i++ {
		if size < 1000 {
			unit = byteUnits[i-1]
			break
		}

		size = size / 1000
	}

	return fmt.Sprintf("%.3g %s", size, unit)
}
<commit_msg>Changes behavior when not printing to a tty<commit_after>package ioprogress

import (
	"fmt"
	"io"
	"os"
	"strings"

	"golang.org/x/crypto/ssh/terminal"
)

// DrawFunc is the callback type for drawing progress.
type DrawFunc func(int64, int64) error

// DrawTextFormatFunc is a callback used by DrawFuncs that draw text in
// order to format the text into some more human friendly format.
type DrawTextFormatFunc func(int64, int64) string

var defaultDrawFunc DrawFunc

func init() {
	defaultDrawFunc = DrawTerminal(os.Stdout)
}

// isTerminal returns True when w is going to a tty, and false otherwise.
func isTerminal(w io.Writer) bool {
	if f, ok := w.(*os.File); ok {
		return terminal.IsTerminal(int(f.Fd()))
	}
	return false
}

// DrawTerminal returns a DrawFunc that draws a progress bar to an io.Writer
// that is assumed to be a terminal (and therefore respects carriage returns).
func DrawTerminal(w io.Writer) DrawFunc {
	return DrawTerminalf(w, func(progress, total int64) string {
		return fmt.Sprintf("%d/%d", progress, total)
	})
}

// DrawTerminalf returns a DrawFunc that draws a progress bar to an io.Writer
// that is formatted with the given formatting function.
func DrawTerminalf(w io.Writer, f DrawTextFormatFunc) DrawFunc {
	var maxLength int

	return func(progress, total int64) error {
		if progress == -1 && total == -1 {
			_, err := fmt.Fprintf(w, "\n")
			return err
		}

		// Make sure we pad it to the max length we've ever drawn so that
		// we don't have trailing characters.
		line := f(progress, total)
		if len(line) < maxLength {
			line = fmt.Sprintf(
				"%s%s",
				line,
				strings.Repeat(" ", maxLength-len(line)))
		}
		maxLength = len(line)

		terminate := "\r"
		if !isTerminal(w) {
			terminate = "\n"
		}
		_, err := fmt.Fprint(w, line+terminate)
		return err
	}
}

var byteUnits = []string{"B", "KB", "MB", "GB", "TB", "PB"}

// DrawTextFormatBytes is a DrawTextFormatFunc that formats the progress
// and total into human-friendly byte formats.
func DrawTextFormatBytes(progress, total int64) string {
	return fmt.Sprintf("%s/%s", ByteUnitStr(progress), ByteUnitStr(total))
}

// DrawTextFormatBar returns a DrawTextFormatFunc that draws a progress
// bar with the given width (in characters). This can be used in conjunction
// with another DrawTextFormatFunc to create a progress bar with bytes, for
// example:
//
//     bar := DrawTextFormatBar(20)
//     func(progress, total int64) string {
//         return fmt.Sprintf(
//           "%s %s",
//           bar(progress, total),
//           DrawTextFormatBytes(progress, total))
//     }
//
func DrawTextFormatBar(width int64) DrawTextFormatFunc {
	return DrawTextFormatBarForW(width, nil)
}

// DrawTextFormatBarForW returns a DrawTextFormatFunc as described in the docs
// for DrawTextFormatBar, however if the io.Writer passed in is not a tty then
// the returned function will always return "".
func DrawTextFormatBarForW(width int64, w io.Writer) DrawTextFormatFunc {
	if w != nil && !isTerminal(w) {
		return func(progress, total int64) string {
			return ""
		}
	}

	width -= 2

	return func(progress, total int64) string {
		current := int64((float64(progress) / float64(total)) * float64(width))
		return fmt.Sprintf(
			"[%s%s]",
			strings.Repeat("=", int(current)),
			strings.Repeat(" ", int(width-current)))
	}
}

// ByteUnitStr pretty prints a number of bytes.
func ByteUnitStr(n int64) string {
	var unit string
	size := float64(n)
	for i := 1; i < len(byteUnits); i++ {
		if size < 1000 {
			unit = byteUnits[i-1]
			break
		}

		size = size / 1000
	}

	return fmt.Sprintf("%.3g %s", size, unit)
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/go.net/websocket"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strings"
)

//constants
const ADDR = ":8080"
const TEMPDIR = "temp/"
const STATICDIR = "template/"

const (
	REGISTER = iota //0
	MESSAGE         //1
	BLOGIN
	BLOGOUT
)

//structs
type User struct {
	Name string
}

type Command struct {
	Type  int
	Value []byte
}

type Connection struct {
	//User who sends
	Usr *User
	//websocket-connection
	Conn *websocket.Conn
	// Buffered channel of outbound messages.
	send chan *Message
	// chatroom in which the user is
	chatRoom *ChatRoom
}

type Message struct {
	Usr *User
	Msg []byte
}

//global vars
//default chatroom
var chatRoom = ChatRoom{
	name:                  "unity is gay",
	history:               make([]*Message, 0),
	broadcast:             make(chan *Message),
	register:              make(chan *Connection),
	unregister:            make(chan *Connection),
	registeredConnections: make(map[*Connection]bool),
}
var server = Server{
	chatRooms:             make([]*ChatRoom, 0),
	register:              make(chan *Connection),
	unregister:            make(chan *Connection),
	registeredConnections: make(map[*Connection]bool),
}

func (c *Connection) Read() {
	for {
		var message string
		err := websocket.Message.Receive(c.Conn, &message)
		if err != nil {
			break
		}

		fmt.Println(message)

		var cmd *Command
		var userName string
		var found bool
		errm := json.Unmarshal([]byte(message), &cmd)
		if errm != nil {

		} else {
			if cmd.Type == REGISTER {
				//check if user name already is taken.
				for d := range server.registeredConnections {
					if d.Usr != nil && d.Usr.Name == string(cmd.Value) {
						found = true
					}
				}
				if found {
					break
				} else {
					c.Usr = &User{Name: string(cmd.Value)}
				}
			}

			if cmd.Type == MESSAGE {
				//check if the message is a private message.
				r, _ := regexp.MatchString("@", string(cmd.Value))
				if r {
					found = false
					userName = strings.Split(string(cmd.Value), " ")[0]
					userName = strings.Split(userName, "@")[1]
					for d := range server.registeredConnections {
						if d.Usr.Name == userName {
							d.send <- &Message{c.Usr, cmd.Value}
							found = true
						}
					}
					if found {
						c.send <- &Message{c.Usr, cmd.Value}
					} else {
						c.send <- &Message{&User{"From server"}, []byte("user not found")}
					}
					//else send the message to everyone
				} else {
					c.chatRoom.broadcast <- &Message{c.Usr, cmd.Value}
				}
			}

			if cmd.Type == BLOGIN {
				c.chatRoom.broadcast <- &Message{c.Usr, []byte("Logged In")}
			}

			if cmd.Type == BLOGOUT {
				c.chatRoom.broadcast <- &Message{c.Usr, []byte("Logged Out")}
				break
			}

		}

	}
	c.Conn.Close()
}

func (c *Connection) Write() {
	for message := range c.send {
		jsonM, _ := json.Marshal(message)
		fmt.Println(string(jsonM))
		err := websocket.Message.Send(c.Conn, string(jsonM))
		if err != nil {
			break
		}
	}
}

func main() {

	go server.run()

	//cleanup if command+c
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		<-c
		server.cleanUp()
		os.Exit(0)
	}()

	http.HandleFunc("/upload", uploadHandler)
	http.HandleFunc("/", mainHandler)
	http.HandleFunc("/index.html", loginHandler)
	http.HandleFunc("/chat.html", chatHandler)
	http.HandleFunc("/doLogin", doLogin)
	http.Handle("/ws", websocket.Handler(wsHandler))
	err := http.ListenAndServe(ADDR, nil)
	if err != nil {
		panic("ListenAndServe: " + err.Error())
	}
}
<commit_msg>import new command package<commit_after>package main

import (
	"code.google.com/p/go.net/websocket"
	"encoding/json"
	"fmt"
	"gofire/command"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strings"
)

//constants
const ADDR = ":8080"
const TEMPDIR = "temp/"
const STATICDIR = "template/"

type User struct {
	Name string
}

type Connection struct {
	//User who sends
	Usr *User
	//websocket-connection
	Conn *websocket.Conn
	// Buffered channel of outbound messages.
	send chan *Message
	// chatroom in which the user is
	chatRoom *ChatRoom
}

type Message struct {
	Usr *User
	Msg []byte
}

//global vars
//default chatroom
var chatRoom = ChatRoom{
	name:                  "unity is gay",
	history:               make([]*Message, 0),
	broadcast:             make(chan *Message),
	register:              make(chan *Connection),
	unregister:            make(chan *Connection),
	registeredConnections: make(map[*Connection]bool),
}
var server = Server{
	chatRooms:             make([]*ChatRoom, 0),
	register:              make(chan *Connection),
	unregister:            make(chan *Connection),
	registeredConnections: make(map[*Connection]bool),
}

func (c *Connection) Read() {
	for {
		var message string
		err := websocket.Message.Receive(c.Conn, &message)
		if err != nil {
			break
		}

		fmt.Println(message)

		var cmd *command.Command
		var userName string
		var found bool
		errm := json.Unmarshal([]byte(message), &cmd)
		if errm != nil {

		} else {
			if cmd.Type == command.REGISTER {
				//check if user name already is taken.
				for d := range server.registeredConnections {
					if d.Usr != nil && d.Usr.Name == string(cmd.Value) {
						found = true
					}
				}
				if found {
					break
				} else {
					c.Usr = &User{Name: string(cmd.Value)}
				}
			}

			if cmd.Type == command.MESSAGE {
				//check if the message is a private message.
				r, _ := regexp.MatchString("@", string(cmd.Value))
				if r {
					found = false
					userName = strings.Split(string(cmd.Value), " ")[0]
					userName = strings.Split(userName, "@")[1]
					for d := range server.registeredConnections {
						if d.Usr.Name == userName {
							d.send <- &Message{c.Usr, cmd.Value}
							found = true
						}
					}
					if found {
						c.send <- &Message{c.Usr, cmd.Value}
					} else {
						c.send <- &Message{&User{"From server"}, []byte("user not found")}
					}
					//else send the message to everyone
				} else {
					c.chatRoom.broadcast <- &Message{c.Usr, cmd.Value}
				}
			}

			if cmd.Type == command.BLOGIN {
				c.chatRoom.broadcast <- &Message{c.Usr, []byte("Logged In")}
			}

			if cmd.Type == command.BLOGOUT {
				c.chatRoom.broadcast <- &Message{c.Usr, []byte("Logged Out")}
				break
			}

		}

	}
	c.Conn.Close()
}

func (c *Connection) Write() {
	for message := range c.send {
		jsonM, _ := json.Marshal(message)
		fmt.Println(string(jsonM))
		err := websocket.Message.Send(c.Conn, string(jsonM))
		if err != nil {
			break
		}
	}
}

func main() {

	go server.run()

	//cleanup if command+c
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		<-c
		server.cleanUp()
		os.Exit(0)
	}()

	http.HandleFunc("/upload", uploadHandler)
	http.HandleFunc("/", mainHandler)
	http.HandleFunc("/index.html", loginHandler)
	http.HandleFunc("/chat.html", chatHandler)
	http.HandleFunc("/doLogin", doLogin)
	http.Handle("/ws", websocket.Handler(wsHandler))
	err := http.ListenAndServe(ADDR, nil)
	if err != nil {
		panic("ListenAndServe: " + err.Error())
	}
}
<|endoftext|>


<commit_before>package nbl

import (
	"sync/atomic"
)

// Int64 implements int64 value for one writer and multiple readers. It is
// designed for architectures that does not support 64-bit atomic operations.
type Int64 struct {
	val [2]int64
	aba uintptr
}

func (u *Int64) StartLoad() uintptr {
	return atomic.LoadUintptr(&u.aba)
}

func (i *Int64) TryLoad(aba uintptr) int64 {
	return i.val[aba&1]
}

func (i *Int64) CheckLoad(aba uintptr) (uintptr, bool) {
	aba1 := atomic.LoadUintptr(&i.aba)
	return aba1, aba1 == aba
}

// BUG: On 32bit CPUs Load doesnt guarantee that it returns valid value. The
// probability of failure depends on the frequency of updates:
// 1 kHz: aba wraps onece per 1193 houres,
// 1 MHz: aba wraps once per 72 minutes.
// Load fails if aba wraps beetwen StartLoad and CheckLoad or between subsequent
// CheckLoads and is read twice with the same value.
func (i *Int64) Load() int64 {
	aba := i.StartLoad()
	for {
		v := i.TryLoad(aba)
		var ok bool
		if aba, ok = i.CheckLoad(aba); ok {
			return v
		}
	}
}

// WriterLoad is more efficient than Load but there should be guarantee that
// any writer can not write to i at the same time.
func (i *Int64) WriterLoad() int64 {
	return i.val[i.aba&1]
}

// WriterStore stores v into i.
// Only ONE writer can call WriterStore at the same time.
func (i *Int64) WriterStore(v int64) {
	aba := i.aba
	aba++
	i.val[aba&1] = v
	atomic.StoreUintptr(&i.aba, aba)
}

// WriterAdd performs i += v and returns new value of i.
// Only ONE writer can call WriterAdd at the same time.
func (i *Int64) WriterAdd(v int64) int64 {
	v += i.WriterLoad()
	i.WriterStore(v)
	return v
}

// WriterSub performs i -= v and returns new value of i.
// Only ONE writer can call WriterSub at the same time.
func (i *Int64) WriterSub(v int64) int64 {
	v = i.WriterLoad() - v
	i.WriterStore(v)
	return v
}
<commit_msg>nbl: Remove unneed Int64.WriterSub method.<commit_after>package nbl

import (
	"sync/atomic"
)

// Int64 implements int64 value for one writer and multiple readers. It is
// designed for architectures that does not support 64-bit atomic operations.
type Int64 struct {
	val [2]int64
	aba uintptr
}

func (u *Int64) StartLoad() uintptr {
	return atomic.LoadUintptr(&u.aba)
}

func (i *Int64) TryLoad(aba uintptr) int64 {
	return i.val[aba&1]
}

func (i *Int64) CheckLoad(aba uintptr) (uintptr, bool) {
	aba1 := atomic.LoadUintptr(&i.aba)
	return aba1, aba1 == aba
}

// BUG: On 32bit CPUs Load does not guarantee that it returns valid value. The
// probability of failure depends on the frequency of updates:
// 1 kHz: aba wraps onece per 1193 houres,
// 1 MHz: aba wraps once per 72 minutes.
// Load fails if aba wraps beetwen StartLoad and CheckLoad or between subsequent
// CheckLoads and after that is readed second time with the same value.
func (i *Int64) Load() int64 {
	aba := i.StartLoad()
	for {
		v := i.TryLoad(aba)
		var ok bool
		if aba, ok = i.CheckLoad(aba); ok {
			return v
		}
	}
}

// WriterLoad is more efficient than Load but there should be guarantee that
// any writer can not write to i at the same time.
func (i *Int64) WriterLoad() int64 {
	return i.val[i.aba&1]
}

// WriterStore stores v into i.
// Only ONE writer can call WriterStore at the same time.
func (i *Int64) WriterStore(v int64) {
	aba := i.aba
	aba++
	i.val[aba&1] = v
	atomic.StoreUintptr(&i.aba, aba)
}

// WriterAdd performs i += v and returns new value of i.
// Only ONE writer can call WriterAdd at the same time.
func (i *Int64) WriterAdd(v int64) int64 {
	v += i.WriterLoad()
	i.WriterStore(v)
	return v
}
<|endoftext|>