<commit_before>package main

import (
	"encoding/base64"
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"strings"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/revert"
	"github.com/lxc/lxd/lxd/util"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
)

func lxdIsConfigured(client lxd.InstanceServer) (bool, error) {
	// Look for networks.
	networks, err := client.GetNetworkNames()
	if err != nil {
		return false, fmt.Errorf("Failed to list networks: %w", err)
	}
	if !shared.StringInSlice("lxdbr0", networks) {
		// Couldn't find lxdbr0.
		return false, nil
	}

	// Look for storage pools.
	pools, err := client.GetStoragePoolNames()
	if err != nil {
		return false, fmt.Errorf("Failed to list storage pools: %w", err)
	}

	if !shared.StringInSlice("default", pools) {
		// No storage pool found.
		return false, nil
	}

	return true, nil
}

func lxdInitialConfiguration(client lxd.InstanceServer) error {
	// Load current server config.
	info, _, err := client.GetServer()
	if err != nil {
		return fmt.Errorf("Failed to get server info: %w", err)
	}
	availableBackends := util.AvailableStorageDrivers(info.Environment.StorageSupportedDrivers, util.PoolTypeLocal)

	// Load the default profile.
	profile, profileEtag, err := client.GetProfile("default")
	if err != nil {
		return fmt.Errorf("Failed to load default profile: %w", err)
	}

	// Look for storage pools.
	pools, err := client.GetStoragePools()
	if err != nil {
		return fmt.Errorf("Failed to list storage pools: %w", err)
	}

	if len(pools) == 0 {
		pool := api.StoragePoolsPost{}
		pool.Config = map[string]string{}
		pool.Name = "default"

		// Check if ZFS supported.
		if shared.StringInSlice("zfs", availableBackends) {
			pool.Driver = "zfs"

			// Check if zsys.
			poolName, _ := shared.RunCommand("zpool", "get", "-H", "-o", "value", "name", "rpool")
			if strings.TrimSpace(poolName) == "rpool" {
				pool.Config["source"] = "rpool/lxd"
			}
		} else {
			// Fallback to dir backend.
			pool.Driver = "dir"
		}

		// Create the storage pool.
		err := client.CreateStoragePool(pool)
		if err != nil {
			return fmt.Errorf("Failed to create storage pool: %w", err)
		}

		// Add to default profile in default project.
		profile.Devices["root"] = map[string]string{
			"type": "disk",
			"pool": "default",
			"path": "/",
		}
	}

	// Look for networks.
	networks, err := client.GetNetworks()
	if err != nil {
		return fmt.Errorf("Failed to list networks: %w", err)
	}

	found := false
	for _, network := range networks {
		if network.Managed {
			found = true
			break
		}
	}

	if !found {
		// Create lxdbr0.
		network := api.NetworksPost{}
		network.Config = map[string]string{}
		network.Type = "bridge"
		network.Name = "lxdbr0"

		err := client.CreateNetwork(network)
		if err != nil {
			return fmt.Errorf("Failed to create network: %w", err)
		}

		// Add to default profile in default project.
		profile.Devices["eth0"] = map[string]string{
			"type":    "nic",
			"network": "default",
			"name":    "eth0",
		}
	}

	// Update the default profile.
	err = client.UpdateProfile("default", profile.Writable(), profileEtag)
	if err != nil {
		return fmt.Errorf("Failed to update default profile: %w", err)
	}

	return nil
}

func lxdSetupUser(uid uint32) error {
	projectName := fmt.Sprintf("user-%d", uid)
	userPath := filepath.Join("users", fmt.Sprintf("%d", uid))

	// User account.
	pw, err := user.LookupId(fmt.Sprintf("%d", uid))
	if err != nil {
		return fmt.Errorf("Failed to retrieve user information: %w", err)
	}

	// Setup reverter.
	revert := revert.New()
	defer revert.Fail()

	// Create certificate directory.
	err = os.MkdirAll(userPath, 0700)
	if err != nil {
		return fmt.Errorf("Failed to create user directory: %w", err)
	}
	revert.Add(func() { os.RemoveAll(userPath) })

	// Generate certificate.
	err = shared.FindOrGenCert(filepath.Join(userPath, "client.crt"), filepath.Join(userPath, "client.key"), true, false)
	if err != nil {
		return fmt.Errorf("Failed to generate user certificate: %w", err)
	}

	// Connect to LXD.
	client, err := lxd.ConnectLXDUnix("", nil)
	if err != nil {
		return fmt.Errorf("Unable to connect to LXD: %w", err)
	}

	client.GetServer()

	// Setup the project (with restrictions).
	projects, err := client.GetProjectNames()
	if err != nil {
		return fmt.Errorf("Unable to retrieve project list: %w", err)
	}

	if !shared.StringInSlice(projectName, projects) {
		// Create the project.
		err := client.CreateProject(api.ProjectsPost{
			Name: projectName,
			ProjectPut: api.ProjectPut{
				Description: fmt.Sprintf("User restricted project for %q (%s)", pw.Username, pw.Uid),
				Config: map[string]string{
					"features.images":               "true",
					"features.networks":             "false",
					"features.profiles":             "true",
					"features.storage.volumes":      "true",
					"restricted":                    "true",
					"restricted.containers.nesting": "allow",
					"restricted.devices.disk":       "allow",
					"restricted.devices.disk.paths": pw.HomeDir,
					"restricted.devices.gpu":        "allow",
					"restricted.idmap.uid":          pw.Uid,
					"restricted.idmap.gid":          pw.Gid,
				},
			},
		})
		if err != nil {
			return fmt.Errorf("Unable to create project: %w", err)
		}

		revert.Add(func() { client.DeleteProject(projectName) })
	}

	// Parse the certificate.
	x509Cert, err := shared.ReadCert(filepath.Join(userPath, "client.crt"))
	if err != nil {
		return fmt.Errorf("Unable to read user certificate: %w", err)
	}

	// Add the certificate to the trust store.
	err = client.CreateCertificate(api.CertificatesPost{
		CertificatePut: api.CertificatePut{
			Name:        fmt.Sprintf("lxd-user-%d", uid),
			Type:        "client",
			Restricted:  true,
			Projects:    []string{projectName},
			Certificate: base64.StdEncoding.EncodeToString(x509Cert.Raw),
		},
	})
	if err != nil {
		return fmt.Errorf("Unable to add user certificate: %w", err)
	}

	revert.Add(func() { client.DeleteCertificate(shared.CertFingerprint(x509Cert)) })

	// Setup default profile.
	err = client.UseProject(projectName).UpdateProfile("default", api.ProfilePut{
		Description: "Default LXD profile",
		Config: map[string]string{
			"raw.idmap": fmt.Sprintf("uid %s %s\ngid %s %s", pw.Uid, pw.Uid, pw.Gid, pw.Gid),
		},
		Devices: map[string]map[string]string{
			"root": {
				"type": "disk",
				"path": "/",
				"pool": "default",
			},
			"eth0": {
				"type":    "nic",
				"name":    "eth0",
				"network": "lxdbr0",
			},
		},
	}, "")
	if err != nil {
		return fmt.Errorf("Unable to update the default profile: %w", err)
	}

	revert.Success()
	return nil
}
<commit_msg>lxd-user: Fix default network name<commit_after>package main

import (
	"encoding/base64"
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"strings"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/revert"
	"github.com/lxc/lxd/lxd/util"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
)

func lxdIsConfigured(client lxd.InstanceServer) (bool, error) {
	// Look for networks.
	networks, err := client.GetNetworkNames()
	if err != nil {
		return false, fmt.Errorf("Failed to list networks: %w", err)
	}
	if !shared.StringInSlice("lxdbr0", networks) {
		// Couldn't find lxdbr0.
		return false, nil
	}

	// Look for storage pools.
	pools, err := client.GetStoragePoolNames()
	if err != nil {
		return false, fmt.Errorf("Failed to list storage pools: %w", err)
	}

	if !shared.StringInSlice("default", pools) {
		// No storage pool found.
		return false, nil
	}

	return true, nil
}

func lxdInitialConfiguration(client lxd.InstanceServer) error {
	// Load current server config.
	info, _, err := client.GetServer()
	if err != nil {
		return fmt.Errorf("Failed to get server info: %w", err)
	}
	availableBackends := util.AvailableStorageDrivers(info.Environment.StorageSupportedDrivers, util.PoolTypeLocal)

	// Load the default profile.
	profile, profileEtag, err := client.GetProfile("default")
	if err != nil {
		return fmt.Errorf("Failed to load default profile: %w", err)
	}

	// Look for storage pools.
	pools, err := client.GetStoragePools()
	if err != nil {
		return fmt.Errorf("Failed to list storage pools: %w", err)
	}

	if len(pools) == 0 {
		pool := api.StoragePoolsPost{}
		pool.Config = map[string]string{}
		pool.Name = "default"

		// Check if ZFS supported.
		if shared.StringInSlice("zfs", availableBackends) {
			pool.Driver = "zfs"

			// Check if zsys.
			poolName, _ := shared.RunCommand("zpool", "get", "-H", "-o", "value", "name", "rpool")
			if strings.TrimSpace(poolName) == "rpool" {
				pool.Config["source"] = "rpool/lxd"
			}
		} else {
			// Fallback to dir backend.
			pool.Driver = "dir"
		}

		// Create the storage pool.
		err := client.CreateStoragePool(pool)
		if err != nil {
			return fmt.Errorf("Failed to create storage pool: %w", err)
		}

		// Add to default profile in default project.
		profile.Devices["root"] = map[string]string{
			"type": "disk",
			"pool": "default",
			"path": "/",
		}
	}

	// Look for networks.
	networks, err := client.GetNetworks()
	if err != nil {
		return fmt.Errorf("Failed to list networks: %w", err)
	}

	found := false
	for _, network := range networks {
		if network.Managed {
			found = true
			break
		}
	}

	if !found {
		// Create lxdbr0.
		network := api.NetworksPost{}
		network.Config = map[string]string{}
		network.Type = "bridge"
		network.Name = "lxdbr0"

		err := client.CreateNetwork(network)
		if err != nil {
			return fmt.Errorf("Failed to create network: %w", err)
		}

		// Add to default profile in default project.
		profile.Devices["eth0"] = map[string]string{
			"type":    "nic",
			"network": "lxdbr0",
			"name":    "eth0",
		}
	}

	// Update the default profile.
	err = client.UpdateProfile("default", profile.Writable(), profileEtag)
	if err != nil {
		return fmt.Errorf("Failed to update default profile: %w", err)
	}

	return nil
}

func lxdSetupUser(uid uint32) error {
	projectName := fmt.Sprintf("user-%d", uid)
	userPath := filepath.Join("users", fmt.Sprintf("%d", uid))

	// User account.
	pw, err := user.LookupId(fmt.Sprintf("%d", uid))
	if err != nil {
		return fmt.Errorf("Failed to retrieve user information: %w", err)
	}

	// Setup reverter.
	revert := revert.New()
	defer revert.Fail()

	// Create certificate directory.
	err = os.MkdirAll(userPath, 0700)
	if err != nil {
		return fmt.Errorf("Failed to create user directory: %w", err)
	}
	revert.Add(func() { os.RemoveAll(userPath) })

	// Generate certificate.
	err = shared.FindOrGenCert(filepath.Join(userPath, "client.crt"), filepath.Join(userPath, "client.key"), true, false)
	if err != nil {
		return fmt.Errorf("Failed to generate user certificate: %w", err)
	}

	// Connect to LXD.
	client, err := lxd.ConnectLXDUnix("", nil)
	if err != nil {
		return fmt.Errorf("Unable to connect to LXD: %w", err)
	}

	client.GetServer()

	// Setup the project (with restrictions).
	projects, err := client.GetProjectNames()
	if err != nil {
		return fmt.Errorf("Unable to retrieve project list: %w", err)
	}

	if !shared.StringInSlice(projectName, projects) {
		// Create the project.
		err := client.CreateProject(api.ProjectsPost{
			Name: projectName,
			ProjectPut: api.ProjectPut{
				Description: fmt.Sprintf("User restricted project for %q (%s)", pw.Username, pw.Uid),
				Config: map[string]string{
					"features.images":               "true",
					"features.networks":             "false",
					"features.profiles":             "true",
					"features.storage.volumes":      "true",
					"restricted":                    "true",
					"restricted.containers.nesting": "allow",
					"restricted.devices.disk":       "allow",
					"restricted.devices.disk.paths": pw.HomeDir,
					"restricted.devices.gpu":        "allow",
					"restricted.idmap.uid":          pw.Uid,
					"restricted.idmap.gid":          pw.Gid,
				},
			},
		})
		if err != nil {
			return fmt.Errorf("Unable to create project: %w", err)
		}

		revert.Add(func() { client.DeleteProject(projectName) })
	}

	// Parse the certificate.
	x509Cert, err := shared.ReadCert(filepath.Join(userPath, "client.crt"))
	if err != nil {
		return fmt.Errorf("Unable to read user certificate: %w", err)
	}

	// Add the certificate to the trust store.
	err = client.CreateCertificate(api.CertificatesPost{
		CertificatePut: api.CertificatePut{
			Name:        fmt.Sprintf("lxd-user-%d", uid),
			Type:        "client",
			Restricted:  true,
			Projects:    []string{projectName},
			Certificate: base64.StdEncoding.EncodeToString(x509Cert.Raw),
		},
	})
	if err != nil {
		return fmt.Errorf("Unable to add user certificate: %w", err)
	}

	revert.Add(func() { client.DeleteCertificate(shared.CertFingerprint(x509Cert)) })

	// Setup default profile.
	err = client.UseProject(projectName).UpdateProfile("default", api.ProfilePut{
		Description: "Default LXD profile",
		Config: map[string]string{
			"raw.idmap": fmt.Sprintf("uid %s %s\ngid %s %s", pw.Uid, pw.Uid, pw.Gid, pw.Gid),
		},
		Devices: map[string]map[string]string{
			"root": {
				"type": "disk",
				"path": "/",
				"pool": "default",
			},
			"eth0": {
				"type":    "nic",
				"name":    "eth0",
				"network": "lxdbr0",
			},
		},
	}, "")
	if err != nil {
		return fmt.Errorf("Unable to update the default profile: %w", err)
	}

	revert.Success()
	return nil
}
<|endoftext|>


<commit_before>package core

import (
	"encoding/base64"
	"errors"
	"fmt"

	ds "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/datastore.go"
	b58 "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-base58"
	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
	"github.com/jbenet/go-ipfs/bitswap"
	bserv "github.com/jbenet/go-ipfs/blockservice"
	config "github.com/jbenet/go-ipfs/config"
	ci "github.com/jbenet/go-ipfs/crypto"
	merkledag "github.com/jbenet/go-ipfs/merkledag"
	path "github.com/jbenet/go-ipfs/path"
	peer "github.com/jbenet/go-ipfs/peer"
	routing "github.com/jbenet/go-ipfs/routing"
	dht "github.com/jbenet/go-ipfs/routing/dht"
	swarm "github.com/jbenet/go-ipfs/swarm"
	u "github.com/jbenet/go-ipfs/util"
)

// IpfsNode is IPFS Core module. It represents an IPFS instance.
type IpfsNode struct {

	// the node's configuration
	Config *config.Config

	// the local node's identity
	Identity *peer.Peer

	// the map of other nodes (Peer instances)
	PeerMap *peer.Map

	// the local datastore
	Datastore ds.Datastore

	// the network message stream
	Swarm *swarm.Swarm

	// the routing system. recommend ipfs-dht
	Routing routing.IpfsRouting

	// the block exchange + strategy (bitswap)
	BitSwap *bitswap.BitSwap

	// the block service, get/add blocks.
	Blocks *bserv.BlockService

	// the merkle dag service, get/add objects.
	DAG *merkledag.DAGService

	// the path resolution system
	Resolver *path.Resolver

	// the name system, resolves paths to hashes
	// Namesys *namesys.Namesys
}

// NewIpfsNode constructs a new IpfsNode based on the given config.
func NewIpfsNode(cfg *config.Config, online bool) (*IpfsNode, error) {
	if cfg == nil {
		return nil, fmt.Errorf("configuration required")
	}

	d, err := makeDatastore(cfg.Datastore)
	if err != nil {
		return nil, err
	}

	local, err := initIdentity(cfg)
	if err != nil {
		return nil, err
	}

	var (
		net *swarm.Swarm
		// TODO: refactor so we can use IpfsRouting interface instead of being DHT-specific
		route* dht.IpfsDHT
		swap *bitswap.BitSwap
	)

	if online {
		net = swarm.NewSwarm(local)
		err = net.Listen()
		if err != nil {
			return nil, err
		}

		route = dht.NewDHT(local, net, d)
		route.Start()

		swap = bitswap.NewBitSwap(local, net, d, route)
		swap.SetStrategy(bitswap.YesManStrategy)

		go initConnections(cfg, route)
	}

	bs, err := bserv.NewBlockService(d, swap)
	if err != nil {
		return nil, err
	}

	dag := &merkledag.DAGService{Blocks: bs}

	return &IpfsNode{
		Config:    cfg,
		PeerMap:   &peer.Map{},
		Datastore: d,
		Blocks:    bs,
		DAG:       dag,
		Resolver:  &path.Resolver{DAG: dag},
		BitSwap:   swap,
		Identity:  local,
		Routing:   route,
	}, nil
}

func initIdentity(cfg *config.Config) (*peer.Peer, error) {
	if len(cfg.Identity.PeerID) == 0 {
		return nil, errors.New("No peer ID in config! (was ipfs init run?)")
	}

	// address is optional
	var addresses []*ma.Multiaddr
	if len(cfg.Identity.Address) > 0 {
		maddr, err := ma.NewMultiaddr(cfg.Identity.Address)
		if err != nil {
			return nil, err
		}

		addresses = []*ma.Multiaddr{ maddr }
	}

	skb, err := base64.StdEncoding.DecodeString(cfg.Identity.PrivKey)
	if err != nil {
		return nil, err
	}

	sk, err := ci.UnmarshalPrivateKey(skb)
	if err != nil {
		return nil, err
	}

	return &peer.Peer{
		ID:        peer.ID(b58.Decode(cfg.Identity.PeerID)),
		Addresses: addresses,
		PrivKey:   sk,
		PubKey:    sk.GetPublic(),
	}, nil
}

func initConnections(cfg *config.Config, route *dht.IpfsDHT) {
	for _, p := range cfg.Peers {
		maddr, err := ma.NewMultiaddr(p.Address)
		if err != nil {
			u.PErr("error: %v\n", err)
			continue
		}

		_, err = route.Connect(maddr)
		if err != nil {
			u.PErr("Bootstrapping error: %v\n", err)
		}
	}
}

func (n *IpfsNode) PinDagNode(nd *merkledag.Node) error {
	u.POut("Pinning node. Currently No-Op\n")
	return nil
}
<commit_msg>core: Show error when config identity is not set<commit_after>package core

import (
	"encoding/base64"
	"errors"
	"fmt"

	ds "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/datastore.go"
	b58 "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-base58"
	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
	"github.com/jbenet/go-ipfs/bitswap"
	bserv "github.com/jbenet/go-ipfs/blockservice"
	config "github.com/jbenet/go-ipfs/config"
	ci "github.com/jbenet/go-ipfs/crypto"
	merkledag "github.com/jbenet/go-ipfs/merkledag"
	path "github.com/jbenet/go-ipfs/path"
	peer "github.com/jbenet/go-ipfs/peer"
	routing "github.com/jbenet/go-ipfs/routing"
	dht "github.com/jbenet/go-ipfs/routing/dht"
	swarm "github.com/jbenet/go-ipfs/swarm"
	u "github.com/jbenet/go-ipfs/util"
)

// IpfsNode is IPFS Core module. It represents an IPFS instance.
type IpfsNode struct {

	// the node's configuration
	Config *config.Config

	// the local node's identity
	Identity *peer.Peer

	// the map of other nodes (Peer instances)
	PeerMap *peer.Map

	// the local datastore
	Datastore ds.Datastore

	// the network message stream
	Swarm *swarm.Swarm

	// the routing system. recommend ipfs-dht
	Routing routing.IpfsRouting

	// the block exchange + strategy (bitswap)
	BitSwap *bitswap.BitSwap

	// the block service, get/add blocks.
	Blocks *bserv.BlockService

	// the merkle dag service, get/add objects.
	DAG *merkledag.DAGService

	// the path resolution system
	Resolver *path.Resolver

	// the name system, resolves paths to hashes
	// Namesys *namesys.Namesys
}

// NewIpfsNode constructs a new IpfsNode based on the given config.
func NewIpfsNode(cfg *config.Config, online bool) (*IpfsNode, error) {
	if cfg == nil {
		return nil, fmt.Errorf("configuration required")
	}

	d, err := makeDatastore(cfg.Datastore)
	if err != nil {
		return nil, err
	}

	local, err := initIdentity(cfg)
	if err != nil {
		return nil, err
	}

	var (
		net *swarm.Swarm
		// TODO: refactor so we can use IpfsRouting interface instead of being DHT-specific
		route* dht.IpfsDHT
		swap *bitswap.BitSwap
	)

	if online {
		net = swarm.NewSwarm(local)
		err = net.Listen()
		if err != nil {
			return nil, err
		}

		route = dht.NewDHT(local, net, d)
		route.Start()

		swap = bitswap.NewBitSwap(local, net, d, route)
		swap.SetStrategy(bitswap.YesManStrategy)

		go initConnections(cfg, route)
	}

	bs, err := bserv.NewBlockService(d, swap)
	if err != nil {
		return nil, err
	}

	dag := &merkledag.DAGService{Blocks: bs}

	return &IpfsNode{
		Config:    cfg,
		PeerMap:   &peer.Map{},
		Datastore: d,
		Blocks:    bs,
		DAG:       dag,
		Resolver:  &path.Resolver{DAG: dag},
		BitSwap:   swap,
		Identity:  local,
		Routing:   route,
	}, nil
}

func initIdentity(cfg *config.Config) (*peer.Peer, error) {
	if cfg.Identity == nil {
		return nil, errors.New("Identity was not set in config (was ipfs init run?)")
	}

	if len(cfg.Identity.PeerID) == 0 {
		return nil, errors.New("No peer ID in config! (was ipfs init run?)")
	}

	// address is optional
	var addresses []*ma.Multiaddr
	if len(cfg.Identity.Address) > 0 {
		maddr, err := ma.NewMultiaddr(cfg.Identity.Address)
		if err != nil {
			return nil, err
		}

		addresses = []*ma.Multiaddr{ maddr }
	}

	skb, err := base64.StdEncoding.DecodeString(cfg.Identity.PrivKey)
	if err != nil {
		return nil, err
	}

	sk, err := ci.UnmarshalPrivateKey(skb)
	if err != nil {
		return nil, err
	}

	return &peer.Peer{
		ID:        peer.ID(b58.Decode(cfg.Identity.PeerID)),
		Addresses: addresses,
		PrivKey:   sk,
		PubKey:    sk.GetPublic(),
	}, nil
}

func initConnections(cfg *config.Config, route *dht.IpfsDHT) {
	for _, p := range cfg.Peers {
		maddr, err := ma.NewMultiaddr(p.Address)
		if err != nil {
			u.PErr("error: %v\n", err)
			continue
		}

		_, err = route.Connect(maddr)
		if err != nil {
			u.PErr("Bootstrapping error: %v\n", err)
		}
	}
}

func (n *IpfsNode) PinDagNode(nd *merkledag.Node) error {
	u.POut("Pinning node. Currently No-Op\n")
	return nil
}
<|endoftext|>


<commit_before>package mtr

import (
	"net"
	"os"
	"time"
)

func Mtr(dest [4]byte, options *TracerouteOptions, c ...chan TracerouteHop) (result TracerouteResult, err error) {
	result.Hops = []TracerouteHop{}
	result.DestAddress = dest
	destAddr := AddressString(dest)
	localAddr := "0.0.0.0"
	ipAddr := net.IPAddr{IP: net.ParseIP(destAddr)}
	pid := os.Getpid() & 0xffff
	timeout := time.Duration(options.TimeoutMs()) * time.Millisecond

	mtrResults := make([]MtrResult, options.MaxHops()+1)

	for snt := 0; snt < options.SntSize(); snt++ {
		retry := 0
		for ttl := 1; ttl < options.MaxHops(); ttl++ {
			hop := TracerouteHop{TTL: ttl, Snt: snt}
			if mtrResults[ttl].TTL == 0 {
				mtrResults[ttl] = MtrResult{TTL: ttl, Host: "???", SuccSum: 0, Success: false, LastTime: time.Duration(0), AllTime: time.Duration(0), BestTime: time.Duration(0), WrstTime: time.Duration(0), AvgTime: time.Duration(0)}
			}
			hopReturn, err := Icmp(localAddr, &ipAddr, ttl, pid, timeout)
			if err != nil || !hopReturn.Success {
				//mtrResults[ttl].Success = false
				hop.Loss = (float32)(snt+1-mtrResults[ttl].SuccSum) / (float32)(snt+1) * 100
				hop.Address = mtrResults[ttl].Host
				hop.AvgTime = mtrResults[ttl].AvgTime
				hop.BestTime = mtrResults[ttl].BestTime
				hop.Host = mtrResults[ttl].Host
				hop.LastTime = mtrResults[ttl].LastTime
				hop.Success = false
				hop.WrstTime = mtrResults[ttl].WrstTime
				notifyMtr(hop, c)
				retry++
				if retry >= options.Retries() {
					break
				}
				continue
			}
			retry = 0
			mtrResults[ttl].SuccSum = mtrResults[ttl].SuccSum + 1
			mtrResults[ttl].Host = hopReturn.Addr
			mtrResults[ttl].LastTime = hopReturn.Elapsed
			if mtrResults[ttl].WrstTime == time.Duration(0) || hopReturn.Elapsed > mtrResults[ttl].WrstTime {
				mtrResults[ttl].WrstTime = hopReturn.Elapsed
			}
			if mtrResults[ttl].BestTime == time.Duration(0) || hopReturn.Elapsed < mtrResults[ttl].BestTime {
				mtrResults[ttl].BestTime = hopReturn.Elapsed
			}
			mtrResults[ttl].AllTime += hopReturn.Elapsed
			mtrResults[ttl].AvgTime = time.Duration((int64)(mtrResults[ttl].AllTime/time.Microsecond)/(int64)(mtrResults[ttl].SuccSum)) * time.Microsecond

			hop.Loss = (float32)(snt+1-mtrResults[ttl].SuccSum) / (float32)(snt+1) * 100
			hop.Address = mtrResults[ttl].Host
			hop.AvgTime = mtrResults[ttl].AvgTime
			hop.BestTime = mtrResults[ttl].BestTime
			hop.Host = mtrResults[ttl].Host
			hop.LastTime = mtrResults[ttl].LastTime
			hop.Success = true
			hop.WrstTime = mtrResults[ttl].WrstTime
			notifyMtr(hop, c)

			if hop.Address == destAddr {
				break
			}
		}
	}

	retry := 0
	for _, mtrResult := range mtrResults {
		if mtrResult.TTL == 0 {
			retry++
			if retry >= options.Retries() {
				break
			}
			continue
		}
		retry = 0
		hop := TracerouteHop{TTL: mtrResult.TTL}
		hop.Address = mtrResult.Host
		hop.Host = mtrResult.Host
		hop.AvgTime = mtrResult.AvgTime
		hop.BestTime = mtrResult.BestTime
		hop.LastTime = mtrResult.LastTime
		failSum := options.SntSize() - mtrResult.SuccSum
		loss := (float32)(failSum) / (float32)(options.SntSize()) * 100
		hop.Loss = float32(loss)
		hop.WrstTime = mtrResult.WrstTime
		hop.Success = true

		result.Hops = append(result.Hops, hop)

		if hop.Host == destAddr {
			break
		}
	}

	closeNotify(c)

	return result, nil
}

func notifyMtr(hop TracerouteHop, channels []chan TracerouteHop) {
	for _, c := range channels {
		c <- hop
	}
}

func closeNotifyMtr(channels []chan TracerouteHop) {
	for _, c := range channels {
		close(c)
	}
}
<commit_msg>0 -> 1<commit_after>package mtr

import (
	"net"
	"os"
	"time"
)

func Mtr(dest [4]byte, options *TracerouteOptions, c ...chan TracerouteHop) (result TracerouteResult, err error) {
	result.Hops = []TracerouteHop{}
	result.DestAddress = dest
	destAddr := AddressString(dest)
	localAddr := "0.0.0.0"
	ipAddr := net.IPAddr{IP: net.ParseIP(destAddr)}
	pid := os.Getpid() & 0xffff
	timeout := time.Duration(options.TimeoutMs()) * time.Millisecond

	mtrResults := make([]MtrResult, options.MaxHops()+1)

	for snt := 0; snt < options.SntSize(); snt++ {
		retry := 0
		for ttl := 1; ttl < options.MaxHops(); ttl++ {
			snt_plus := snt + 1
			hop := TracerouteHop{TTL: ttl, Snt: snt_plus}
			if mtrResults[ttl].TTL == 0 {
				mtrResults[ttl] = MtrResult{TTL: ttl, Host: "???", SuccSum: 0, Success: false, LastTime: time.Duration(0), AllTime: time.Duration(0), BestTime: time.Duration(0), WrstTime: time.Duration(0), AvgTime: time.Duration(0)}
			}
			hopReturn, err := Icmp(localAddr, &ipAddr, ttl, pid, timeout)
			if err != nil || !hopReturn.Success {
				//mtrResults[ttl].Success = false
				hop.Loss = (float32)(snt+1-mtrResults[ttl].SuccSum) / (float32)(snt+1) * 100
				hop.Address = mtrResults[ttl].Host
				hop.AvgTime = mtrResults[ttl].AvgTime
				hop.BestTime = mtrResults[ttl].BestTime
				hop.Host = mtrResults[ttl].Host
				hop.LastTime = mtrResults[ttl].LastTime
				hop.Success = false
				hop.WrstTime = mtrResults[ttl].WrstTime
				notifyMtr(hop, c)
				retry++
				if retry >= options.Retries() {
					break
				}
				continue
			}
			retry = 0
			mtrResults[ttl].SuccSum = mtrResults[ttl].SuccSum + 1
			mtrResults[ttl].Host = hopReturn.Addr
			mtrResults[ttl].LastTime = hopReturn.Elapsed
			if mtrResults[ttl].WrstTime == time.Duration(0) || hopReturn.Elapsed > mtrResults[ttl].WrstTime {
				mtrResults[ttl].WrstTime = hopReturn.Elapsed
			}
			if mtrResults[ttl].BestTime == time.Duration(0) || hopReturn.Elapsed < mtrResults[ttl].BestTime {
				mtrResults[ttl].BestTime = hopReturn.Elapsed
			}
			mtrResults[ttl].AllTime += hopReturn.Elapsed
			mtrResults[ttl].AvgTime = time.Duration((int64)(mtrResults[ttl].AllTime/time.Microsecond)/(int64)(mtrResults[ttl].SuccSum)) * time.Microsecond

			hop.Loss = (float32)(snt+1-mtrResults[ttl].SuccSum) / (float32)(snt+1) * 100
			hop.Address = mtrResults[ttl].Host
			hop.AvgTime = mtrResults[ttl].AvgTime
			hop.BestTime = mtrResults[ttl].BestTime
			hop.Host = mtrResults[ttl].Host
			hop.LastTime = mtrResults[ttl].LastTime
			hop.Success = true
			hop.WrstTime = mtrResults[ttl].WrstTime
			notifyMtr(hop, c)

			if hop.Address == destAddr {
				break
			}
		}
	}

	retry := 0
	for _, mtrResult := range mtrResults {
		if mtrResult.TTL == 0 {
			retry++
			if retry >= options.Retries() {
				break
			}
			continue
		}
		retry = 0
		hop := TracerouteHop{TTL: mtrResult.TTL}
		hop.Address = mtrResult.Host
		hop.Host = mtrResult.Host
		hop.AvgTime = mtrResult.AvgTime
		hop.BestTime = mtrResult.BestTime
		hop.LastTime = mtrResult.LastTime
		failSum := options.SntSize() - mtrResult.SuccSum
		loss := (float32)(failSum) / (float32)(options.SntSize()) * 100
		hop.Loss = float32(loss)
		hop.WrstTime = mtrResult.WrstTime
		hop.Success = true

		result.Hops = append(result.Hops, hop)

		if hop.Host == destAddr {
			break
		}
	}

	closeNotify(c)

	return result, nil
}

func notifyMtr(hop TracerouteHop, channels []chan TracerouteHop) {
	for _, c := range channels {
		c <- hop
	}
}

func closeNotifyMtr(channels []chan TracerouteHop) {
	for _, c := range channels {
		close(c)
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 Space Monkey, Inc.
//
// 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 openssl

import (
	"errors"
	"net"
)

type listener struct {
	net.Listener
	ctx *Ctx
}

func (l *listener) Accept() (c net.Conn, err error) {
	c, err = l.Listener.Accept()
	if err != nil {
		return nil, err
	}
	ssl_c, err := Server(c, l.ctx)
	if err != nil {
		c.Close()
		return nil, err
	}
	return ssl_c, nil
}

// NewListener wraps an existing net.Listener such that all accepted
// connections are wrapped as OpenSSL server connections using the provided
// context ctx.
func NewListener(inner net.Listener, ctx *Ctx) net.Listener {
	return &listener{
		Listener: inner,
		ctx:      ctx}
}

// Listen is a wrapper around net.Listen that wraps incoming connections with
// an OpenSSL server connection using the provided context ctx.
func Listen(network, laddr string, ctx *Ctx) (net.Listener, error) {
	if ctx == nil {
		return nil, errors.New("no ssl context provided")
	}
	l, err := net.Listen(network, laddr)
	if err != nil {
		return nil, err
	}
	return NewListener(l, ctx), nil
}

type DialFlags int

const (
	InsecureSkipHostVerification DialFlags = 1 << iota
	DisableSNI
)

// Dial will connect to network/address and then wrap the corresponding
// underlying connection with an OpenSSL client connection using context ctx.
// If flags includes InsecureSkipHostVerification, the server certificate's
// hostname will not be checked to match the hostname in addr. Otherwise, flags
// should be 0.
//
// Dial probably won't work for you unless you set a verify location or add
// some certs to the certificate store of the client context you're using.
// This library is not nice enough to use the system certificate store by
// default for you yet.
func Dial(network, addr string, ctx *Ctx, flags DialFlags) (*Conn, error) {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, err
	}
	if ctx == nil {
		var err error
		ctx, err = NewCtx()
		if err != nil {
			return nil, err
		}
		// TODO: use operating system default certificate chain?
	}
	c, err := net.Dial(network, addr)
	if err != nil {
		return nil, err
	}
	conn, err := Client(c, ctx)
	if err != nil {
		c.Close()
		return nil, err
	}
	if flags&DisableSNI == 0 {
		err = conn.SetTlsExtHostName(host)
		if err != nil {
			conn.Close()
			return nil, err
		}
	}
	err = conn.Handshake()
	if err != nil {
		conn.Close()
		return nil, err
	}
	if flags&InsecureSkipHostVerification == 0 {
	}
	return conn, nil
}
<commit_msg>adding wrapper function which allows customized dialer object<commit_after>// Copyright (C) 2014 Space Monkey, Inc.
//
// 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 openssl

import (
	"errors"
	"net"
)

type listener struct {
	net.Listener
	ctx *Ctx
}

func (l *listener) Accept() (c net.Conn, err error) {
	c, err = l.Listener.Accept()
	if err != nil {
		return nil, err
	}
	ssl_c, err := Server(c, l.ctx)
	if err != nil {
		c.Close()
		return nil, err
	}
	return ssl_c, nil
}

// NewListener wraps an existing net.Listener such that all accepted
// connections are wrapped as OpenSSL server connections using the provided
// context ctx.
func NewListener(inner net.Listener, ctx *Ctx) net.Listener {
	return &listener{
		Listener: inner,
		ctx:      ctx}
}

// Listen is a wrapper around net.Listen that wraps incoming connections with
// an OpenSSL server connection using the provided context ctx.
func Listen(network, laddr string, ctx *Ctx) (net.Listener, error) {
	if ctx == nil {
		return nil, errors.New("no ssl context provided")
	}
	l, err := net.Listen(network, laddr)
	if err != nil {
		return nil, err
	}
	return NewListener(l, ctx), nil
}

type DialFlags int

const (
	InsecureSkipHostVerification DialFlags = 1 << iota
	DisableSNI
)

// Dial will connect to network/address and then wrap the corresponding
// underlying connection with an OpenSSL client connection using context ctx.
// If flags includes InsecureSkipHostVerification, the server certificate's
// hostname will not be checked to match the hostname in addr. Otherwise, flags
// should be 0.
//
// Dial probably won't work for you unless you set a verify location or add
// some certs to the certificate store of the client context you're using.
// This library is not nice enough to use the system certificate store by
// default for you yet.
func Dial(network, addr string, ctx *Ctx, flags DialFlags) (*Conn, error) {
	return DialWithDialer(network, addr, ctx, flags, &net.Dialer{})
}

// DialWithDialer is a wrapper function which allows an additional parameter for a customized dialer objects
func DialWithDialer(network, addr string, ctx *Ctx, flags DialFlags, dialer *net.Dialer) (*Conn, error) {
	if dialer == nil {
		dialer = &net.Dialer{}
	}
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, err
	}
	if ctx == nil {
		var err error
		ctx, err = NewCtx()
		if err != nil {
			return nil, err
		}
		// TODO: use operating system default certificate chain?
	}
	c, err := dialer.Dial(network, addr)
	if err != nil {
		return nil, err
	}
	conn, err := Client(c, ctx)
	if err != nil {
		c.Close()
		return nil, err
	}
	if flags&DisableSNI == 0 {
		err = conn.SetTlsExtHostName(host)
		if err != nil {
			conn.Close()
			return nil, err
		}
	}
	err = conn.Handshake()
	if err != nil {
		conn.Close()
		return nil, err
	}
	if flags&InsecureSkipHostVerification == 0 {
	}
	return conn, nil
}
<|endoftext|>


<commit_before>package nsf

import (
	"errors"
	"io"
	"io/ioutil"
	"time"

	"github.com/mjibson/mog/codec/nsf/cpu6502"
)

const (
	// 1.79 MHz
	cpuClock = 236250000 / 11 / 12
)

var (
	// DefaultSampleRate is the default sample rate of a track after calling
	// Init().
	DefaultSampleRate = 44100
	ErrUnrecognized   = errors.New("nsf: unrecognized format")
)

const (
	NSF_HEADER_LEN = 0x80
	NSF_VERSION    = 0x5
	NSF_SONGS      = 0x6
	NSF_START      = 0x7
	NSF_LOAD       = 0x8
	NSF_INIT       = 0xa
	NSF_PLAY       = 0xc
	NSF_SONG       = 0xe
	NSF_ARTIST     = 0x2e
	NSF_COPYRIGHT  = 0x4e
	NSF_SPEED_NTSC = 0x6e
	NSF_BANKSWITCH = 0x70
	NSF_SPEED_PAL  = 0x78
	NSF_PAL_NTSC   = 0x7a
	NSF_EXTRA      = 0x7b
	NSF_ZERO       = 0x7c
)

func ReadNSF(r io.Reader) (n *NSF, err error) {
	n = &NSF{}
	n.b, err = ioutil.ReadAll(r)
	if err != nil {
		return
	}
	if len(n.b) < NSF_HEADER_LEN ||
		string(n.b[0:NSF_VERSION]) != "NESM\u001a" {
		return nil, ErrUnrecognized
	}
	n.Version = n.b[NSF_VERSION]
	n.Songs = n.b[NSF_SONGS]
	n.Start = n.b[NSF_START]
	n.LoadAddr = bLEtoUint16(n.b[NSF_LOAD:])
	n.InitAddr = bLEtoUint16(n.b[NSF_INIT:])
	n.PlayAddr = bLEtoUint16(n.b[NSF_PLAY:])
	n.Song = bToString(n.b[NSF_SONG:])
	n.Artist = bToString(n.b[NSF_ARTIST:])
	n.Copyright = bToString(n.b[NSF_COPYRIGHT:])
	n.SpeedNTSC = bLEtoUint16(n.b[NSF_SPEED_NTSC:])
	copy(n.Bankswitch[:], n.b[NSF_BANKSWITCH:NSF_SPEED_PAL])
	n.SpeedPAL = bLEtoUint16(n.b[NSF_SPEED_PAL:])
	n.PALNTSC = n.b[NSF_PAL_NTSC]
	n.Extra = n.b[NSF_EXTRA]
	n.Data = n.b[NSF_HEADER_LEN:]
	return
}

type NSF struct {
	*Ram
	*cpu6502.Cpu

	b []byte // raw NSF data

	Version byte
	Songs   byte
	Start   byte

	LoadAddr uint16
	InitAddr uint16
	PlayAddr uint16

	Song      string
	Artist    string
	Copyright string

	SpeedNTSC  uint16
	Bankswitch [8]byte
	SpeedPAL   uint16
	PALNTSC    byte
	Extra      byte
	Data       []byte

	// SampleRate is the sample rate at which samples will be generated. If not
	// set before Init(), it is set to DefaultSampleRate.
	SampleRate  int64
	totalTicks  int64
	frameTicks  int64
	sampleTicks int64
	playTicks   int64
	samples     []float32
	prevs       [4]float32
	pi          int // prevs index
}

func New() *NSF {
	n := NSF{
		Ram: new(Ram),
	}
	n.Cpu = cpu6502.New(n.Ram)
	n.Cpu.DisableDecimal = true
	n.Cpu.P = 0x24
	n.Cpu.S = 0xfd
	return &n
}

func (n *NSF) Tick() {
	n.Ram.A.Step()
	n.totalTicks++
	n.frameTicks++
	if n.frameTicks == cpuClock/240 {
		n.frameTicks = 0
		n.Ram.A.FrameStep()
	}
	n.sampleTicks++
	if n.sampleTicks >= cpuClock/n.SampleRate {
		n.sampleTicks = 0
		n.append(n.Ram.A.Volume())
	}
	n.playTicks++
}

func (n *NSF) append(v float32) {
	n.prevs[n.pi] = v
	n.pi++
	if n.pi >= len(n.prevs) {
		n.pi = 0
	}
	var sum float32
	for _, s := range n.prevs {
		sum += s
	}
	sum /= float32(len(n.prevs))
	n.samples = append(n.samples, sum)
}

func (n *NSF) Init(song byte) {
	if n.SampleRate == 0 {
		n.SampleRate = int64(DefaultSampleRate)
	}
	copy(n.Ram.M[n.LoadAddr:], n.Data)
	n.Ram.A.Init()
	n.Cpu.A = song - 1
	n.Cpu.PC = n.InitAddr
	n.Cpu.T = nil
	n.Cpu.Run()
	n.Cpu.T = n
}

func (n *NSF) Play(samples int) []float32 {
	playDur := time.Duration(n.SpeedNTSC) * time.Nanosecond * 1000
	ticksPerPlay := int64(playDur / (time.Second / cpuClock))
	n.samples = make([]float32, 0, samples)
	for len(n.samples) < samples {
		n.playTicks = 0
		n.Cpu.PC = n.PlayAddr
		n.Cpu.Halt = false
		for !n.Cpu.Halt && len(n.samples) < samples {
			n.Cpu.Step()
			if !n.Cpu.I() {
				panic("INTERRUPT")
			}
		}
		for i := ticksPerPlay - n.playTicks; i > 0 && len(n.samples) < samples; i-- {
			n.Tick()
		}
	}
	return n.samples
}

// little-endian [2]byte to uint16 conversion
func bLEtoUint16(b []byte) uint16 {
	return uint16(b[1])<<8 + uint16(b[0])
}

// null-terminated bytes to string
func bToString(b []byte) string {
	i := 0
	for i = range b {
		if b[i] == 0 {
			break
		}
	}
	return string(b[:i])
}

type Ram struct {
	M [0xffff + 1]byte
	A Apu
}

func (r *Ram) Read(v uint16) byte {
	switch v {
	case 0x4015:
		return r.A.Read(v)
	default:
		return r.M[v]
	}
}

func (r *Ram) Write(v uint16, b byte) {
	r.M[v] = b
	if v&0xf000 == 0x4000 {
		r.A.Write(v, b)
	}
}
<commit_msg>Use a standard Step function, useful for tests<commit_after>package nsf

import (
	"errors"
	"io"
	"io/ioutil"
	"time"

	"github.com/mjibson/mog/codec/nsf/cpu6502"
)

const (
	// 1.79 MHz
	cpuClock = 236250000 / 11 / 12
)

var (
	// DefaultSampleRate is the default sample rate of a track after calling
	// Init().
	DefaultSampleRate = 44100
	ErrUnrecognized   = errors.New("nsf: unrecognized format")
)

const (
	NSF_HEADER_LEN = 0x80
	NSF_VERSION    = 0x5
	NSF_SONGS      = 0x6
	NSF_START      = 0x7
	NSF_LOAD       = 0x8
	NSF_INIT       = 0xa
	NSF_PLAY       = 0xc
	NSF_SONG       = 0xe
	NSF_ARTIST     = 0x2e
	NSF_COPYRIGHT  = 0x4e
	NSF_SPEED_NTSC = 0x6e
	NSF_BANKSWITCH = 0x70
	NSF_SPEED_PAL  = 0x78
	NSF_PAL_NTSC   = 0x7a
	NSF_EXTRA      = 0x7b
	NSF_ZERO       = 0x7c
)

func ReadNSF(r io.Reader) (n *NSF, err error) {
	n = &NSF{}
	n.b, err = ioutil.ReadAll(r)
	if err != nil {
		return
	}
	if len(n.b) < NSF_HEADER_LEN ||
		string(n.b[0:NSF_VERSION]) != "NESM\u001a" {
		return nil, ErrUnrecognized
	}
	n.Version = n.b[NSF_VERSION]
	n.Songs = n.b[NSF_SONGS]
	n.Start = n.b[NSF_START]
	n.LoadAddr = bLEtoUint16(n.b[NSF_LOAD:])
	n.InitAddr = bLEtoUint16(n.b[NSF_INIT:])
	n.PlayAddr = bLEtoUint16(n.b[NSF_PLAY:])
	n.Song = bToString(n.b[NSF_SONG:])
	n.Artist = bToString(n.b[NSF_ARTIST:])
	n.Copyright = bToString(n.b[NSF_COPYRIGHT:])
	n.SpeedNTSC = bLEtoUint16(n.b[NSF_SPEED_NTSC:])
	copy(n.Bankswitch[:], n.b[NSF_BANKSWITCH:NSF_SPEED_PAL])
	n.SpeedPAL = bLEtoUint16(n.b[NSF_SPEED_PAL:])
	n.PALNTSC = n.b[NSF_PAL_NTSC]
	n.Extra = n.b[NSF_EXTRA]
	n.Data = n.b[NSF_HEADER_LEN:]
	return
}

type NSF struct {
	*Ram
	*cpu6502.Cpu

	b []byte // raw NSF data

	Version byte
	Songs   byte
	Start   byte

	LoadAddr uint16
	InitAddr uint16
	PlayAddr uint16

	Song      string
	Artist    string
	Copyright string

	SpeedNTSC  uint16
	Bankswitch [8]byte
	SpeedPAL   uint16
	PALNTSC    byte
	Extra      byte
	Data       []byte

	// SampleRate is the sample rate at which samples will be generated. If not
	// set before Init(), it is set to DefaultSampleRate.
	SampleRate  int64
	totalTicks  int64
	frameTicks  int64
	sampleTicks int64
	playTicks   int64
	samples     []float32
	prevs       [4]float32
	pi          int // prevs index
}

func New() *NSF {
	n := NSF{
		Ram: new(Ram),
	}
	n.Cpu = cpu6502.New(n.Ram)
	n.Cpu.DisableDecimal = true
	n.Cpu.P = 0x24
	n.Cpu.S = 0xfd
	return &n
}

func (n *NSF) Tick() {
	n.Ram.A.Step()
	n.totalTicks++
	n.frameTicks++
	if n.frameTicks == cpuClock/240 {
		n.frameTicks = 0
		n.Ram.A.FrameStep()
	}
	n.sampleTicks++
	if n.sampleTicks >= cpuClock/n.SampleRate {
		n.sampleTicks = 0
		n.append(n.Ram.A.Volume())
	}
	n.playTicks++
}

func (n *NSF) append(v float32) {
	n.prevs[n.pi] = v
	n.pi++
	if n.pi >= len(n.prevs) {
		n.pi = 0
	}
	var sum float32
	for _, s := range n.prevs {
		sum += s
	}
	sum /= float32(len(n.prevs))
	n.samples = append(n.samples, sum)
}

func (n *NSF) Init(song byte) {
	if n.SampleRate == 0 {
		n.SampleRate = int64(DefaultSampleRate)
	}
	copy(n.Ram.M[n.LoadAddr:], n.Data)
	n.Ram.A.Init()
	n.Cpu.A = song - 1
	n.Cpu.PC = n.InitAddr
	n.Cpu.T = nil
	n.Cpu.Run()
	n.Cpu.T = n
}

func (n *NSF) Step() {
	n.Cpu.Step()
	if !n.Cpu.I() && n.Ram.A.Interrupt {
		println("INTERRUPT")
		n.Cpu.Interrupt()
	}
}

func (n *NSF) Play(samples int) []float32 {
	playDur := time.Duration(n.SpeedNTSC) * time.Nanosecond * 1000
	ticksPerPlay := int64(playDur / (time.Second / cpuClock))
	n.samples = make([]float32, 0, samples)
	for len(n.samples) < samples {
		n.playTicks = 0
		n.Cpu.PC = n.PlayAddr
		n.Cpu.Halt = false
		for !n.Cpu.Halt && len(n.samples) < samples {
			n.Step()
		}
		for i := ticksPerPlay - n.playTicks; i > 0 && len(n.samples) < samples; i-- {
			n.Tick()
		}
	}
	return n.samples
}

// little-endian [2]byte to uint16 conversion
func bLEtoUint16(b []byte) uint16 {
	return uint16(b[1])<<8 + uint16(b[0])
}

// null-terminated bytes to string
func bToString(b []byte) string {
	i := 0
	for i = range b {
		if b[i] == 0 {
			break
		}
	}
	return string(b[:i])
}

type Ram struct {
	M [0xffff + 1]byte
	A Apu
}

func (r *Ram) Read(v uint16) byte {
	switch v {
	case 0x4015:
		return r.A.Read(v)
	default:
		return r.M[v]
	}
}

func (r *Ram) Write(v uint16, b byte) {
	r.M[v] = b
	if v&0xf000 == 0x4000 {
		r.A.Write(v, b)
	}
}
<|endoftext|>


<commit_before>package nsf

import (
	"errors"
	"io"
	"io/ioutil"
	"time"

	"github.com/mjibson/mog/codec/nsf/cpu6502"
)

const (
	// 1.79 MHz
	cpuClock          = 236250000 / 11 / 12
	DefaultSampleRate = 44100
)

var (
	ErrUnrecognized = errors.New("nsf: unrecognized format")
)

const (
	NSF_HEADER_LEN = 0x80
	NSF_VERSION    = 0x5
	NSF_SONGS      = 0x6
	NSF_START      = 0x7
	NSF_LOAD       = 0x8
	NSF_INIT       = 0xa
	NSF_PLAY       = 0xc
	NSF_SONG       = 0xe
	NSF_ARTIST     = 0x2e
	NSF_COPYRIGHT  = 0x4e
	NSF_SPEED_NTSC = 0x6e
	NSF_BANKSWITCH = 0x70
	NSF_SPEED_PAL  = 0x78
	NSF_PAL_NTSC   = 0x7a
	NSF_EXTRA      = 0x7b
	NSF_ZERO       = 0x7c
)

func ReadNSF(r io.Reader) (n *NSF, err error) {
	n = &NSF{}
	n.b, err = ioutil.ReadAll(r)
	if err != nil {
		return
	}
	if len(n.b) < NSF_HEADER_LEN ||
		string(n.b[0:NSF_VERSION]) != "NESM\u001a" {
		return nil, ErrUnrecognized
	}
	n.Version = n.b[NSF_VERSION]
	n.Songs = n.b[NSF_SONGS]
	n.Start = n.b[NSF_START]
	n.LoadAddr = bLEtoUint16(n.b[NSF_LOAD:])
	n.InitAddr = bLEtoUint16(n.b[NSF_INIT:])
	n.PlayAddr = bLEtoUint16(n.b[NSF_PLAY:])
	n.Song = bToString(n.b[NSF_SONG:])
	n.Artist = bToString(n.b[NSF_ARTIST:])
	n.Copyright = bToString(n.b[NSF_COPYRIGHT:])
	n.SpeedNTSC = bLEtoUint16(n.b[NSF_SPEED_NTSC:])
	copy(n.Bankswitch[:], n.b[NSF_BANKSWITCH:NSF_SPEED_PAL])
	n.SpeedPAL = bLEtoUint16(n.b[NSF_SPEED_PAL:])
	n.PALNTSC = n.b[NSF_PAL_NTSC]
	n.Extra = n.b[NSF_EXTRA]
	n.Data = n.b[NSF_HEADER_LEN:]
	return
}

type NSF struct {
	*Ram
	*cpu6502.Cpu

	b []byte // raw NSF data

	Version byte
	Songs   byte
	Start   byte

	LoadAddr uint16
	InitAddr uint16
	PlayAddr uint16

	Song      string
	Artist    string
	Copyright string

	SpeedNTSC  uint16
	Bankswitch [8]byte
	SpeedPAL   uint16
	PALNTSC    byte
	Extra      byte
	Data       []byte

	SampleRate  int64
	totalTicks  int64
	frameTicks  int64
	sampleTicks int64
	playTicks   int64
	samples     []float32
	prevs       [4]float32
	pi          int // prevs index
}

func (n *NSF) Tick() {
	n.Ram.A.Step()
	n.totalTicks++
	n.frameTicks++
	if n.frameTicks == cpuClock/240 {
		n.frameTicks = 0
		n.Ram.A.FrameStep()
	}
	n.sampleTicks++
	if n.sampleTicks >= cpuClock/n.SampleRate {
		n.sampleTicks = 0
		n.append(n.Ram.A.Volume())
	}
	n.playTicks++
}

func (n *NSF) append(v float32) {
	n.prevs[n.pi] = v
	n.pi++
	if n.pi >= len(n.prevs) {
		n.pi = 0
	}
	var sum float32
	for _, s := range n.prevs {
		sum += s
	}
	sum /= float32(len(n.prevs))
	n.samples = append(n.samples, sum)
}

func (n *NSF) Init(song byte) {
	if n.SampleRate == 0 {
		n.SampleRate = DefaultSampleRate
	}
	n.Ram = new(Ram)
	n.Cpu = cpu6502.New(n.Ram)
	copy(n.Ram.M[n.LoadAddr:], n.Data)
	n.Ram.A.Init()
	n.Cpu.A = song - 1
	n.Cpu.PC = n.InitAddr
	n.Cpu.T = nil
	n.Cpu.Run()
	n.Cpu.T = n
}

func (n *NSF) Play(samples int) []float32 {
	playDur := time.Duration(n.SpeedNTSC) * time.Nanosecond * 1000
	ticksPerPlay := int64(playDur / (time.Second / cpuClock))
	n.samples = make([]float32, 0, samples)
	for len(n.samples) < samples {
		n.playTicks = 0
		n.Cpu.PC = n.PlayAddr
		n.Cpu.Halt = false
		for !n.Cpu.Halt && len(n.samples) < samples {
			n.Cpu.Step()
			if !n.Cpu.I() {
				panic("INTERRUPT")
			}
		}
		for i := ticksPerPlay - n.playTicks; i > 0 && len(n.samples) < samples; i-- {
			n.Tick()
		}
	}
	return n.samples
}

// little-endian [2]byte to uint16 conversion
func bLEtoUint16(b []byte) uint16 {
	return uint16(b[1])<<8 + uint16(b[0])
}

// null-terminated bytes to string
func bToString(b []byte) string {
	i := 0
	for i = range b {
		if b[i] == 0 {
			break
		}
	}
	return string(b[:i])
}

type Ram struct {
	M [0xffff + 1]byte
	A Apu
}

func (r *Ram) Read(v uint16) byte {
	if v == 0x4015 {
		return r.A.Read(v)
	} else {
		return r.M[v]
	}
}

func (r *Ram) Write(v uint16, b byte) {
	r.M[v] = b
	if v&0xf000 == 0x4000 {
		r.A.Write(v, b)
	}
}
<commit_msg>Document sample rates<commit_after>package nsf

import (
	"errors"
	"io"
	"io/ioutil"
	"time"

	"github.com/mjibson/mog/codec/nsf/cpu6502"
)

const (
	// 1.79 MHz
	cpuClock = 236250000 / 11 / 12
)

var (
	// DefaultSampleRate is the default sample rate of a track after calling
	// Init().
	DefaultSampleRate = 44100
	ErrUnrecognized   = errors.New("nsf: unrecognized format")
)

const (
	NSF_HEADER_LEN = 0x80
	NSF_VERSION    = 0x5
	NSF_SONGS      = 0x6
	NSF_START      = 0x7
	NSF_LOAD       = 0x8
	NSF_INIT       = 0xa
	NSF_PLAY       = 0xc
	NSF_SONG       = 0xe
	NSF_ARTIST     = 0x2e
	NSF_COPYRIGHT  = 0x4e
	NSF_SPEED_NTSC = 0x6e
	NSF_BANKSWITCH = 0x70
	NSF_SPEED_PAL  = 0x78
	NSF_PAL_NTSC   = 0x7a
	NSF_EXTRA      = 0x7b
	NSF_ZERO       = 0x7c
)

func ReadNSF(r io.Reader) (n *NSF, err error) {
	n = &NSF{}
	n.b, err = ioutil.ReadAll(r)
	if err != nil {
		return
	}
	if len(n.b) < NSF_HEADER_LEN ||
		string(n.b[0:NSF_VERSION]) != "NESM\u001a" {
		return nil, ErrUnrecognized
	}
	n.Version = n.b[NSF_VERSION]
	n.Songs = n.b[NSF_SONGS]
	n.Start = n.b[NSF_START]
	n.LoadAddr = bLEtoUint16(n.b[NSF_LOAD:])
	n.InitAddr = bLEtoUint16(n.b[NSF_INIT:])
	n.PlayAddr = bLEtoUint16(n.b[NSF_PLAY:])
	n.Song = bToString(n.b[NSF_SONG:])
	n.Artist = bToString(n.b[NSF_ARTIST:])
	n.Copyright = bToString(n.b[NSF_COPYRIGHT:])
	n.SpeedNTSC = bLEtoUint16(n.b[NSF_SPEED_NTSC:])
	copy(n.Bankswitch[:], n.b[NSF_BANKSWITCH:NSF_SPEED_PAL])
	n.SpeedPAL = bLEtoUint16(n.b[NSF_SPEED_PAL:])
	n.PALNTSC = n.b[NSF_PAL_NTSC]
	n.Extra = n.b[NSF_EXTRA]
	n.Data = n.b[NSF_HEADER_LEN:]
	return
}

type NSF struct {
	*Ram
	*cpu6502.Cpu

	b []byte // raw NSF data

	Version byte
	Songs   byte
	Start   byte

	LoadAddr uint16
	InitAddr uint16
	PlayAddr uint16

	Song      string
	Artist    string
	Copyright string

	SpeedNTSC  uint16
	Bankswitch [8]byte
	SpeedPAL   uint16
	PALNTSC    byte
	Extra      byte
	Data       []byte

	// SampleRate is the sample rate at which samples will be generated. If not
	// set before Init(), it is set to DefaultSampleRate.
	SampleRate  int64
	totalTicks  int64
	frameTicks  int64
	sampleTicks int64
	playTicks   int64
	samples     []float32
	prevs       [4]float32
	pi          int // prevs index
}

func (n *NSF) Tick() {
	n.Ram.A.Step()
	n.totalTicks++
	n.frameTicks++
	if n.frameTicks == cpuClock/240 {
		n.frameTicks = 0
		n.Ram.A.FrameStep()
	}
	n.sampleTicks++
	if n.sampleTicks >= cpuClock/n.SampleRate {
		n.sampleTicks = 0
		n.append(n.Ram.A.Volume())
	}
	n.playTicks++
}

func (n *NSF) append(v float32) {
	n.prevs[n.pi] = v
	n.pi++
	if n.pi >= len(n.prevs) {
		n.pi = 0
	}
	var sum float32
	for _, s := range n.prevs {
		sum += s
	}
	sum /= float32(len(n.prevs))
	n.samples = append(n.samples, sum)
}

func (n *NSF) Init(song byte) {
	if n.SampleRate == 0 {
		n.SampleRate = int64(DefaultSampleRate)
	}
	n.Ram = new(Ram)
	n.Cpu = cpu6502.New(n.Ram)
	copy(n.Ram.M[n.LoadAddr:], n.Data)
	n.Ram.A.Init()
	n.Cpu.A = song - 1
	n.Cpu.PC = n.InitAddr
	n.Cpu.T = nil
	n.Cpu.Run()
	n.Cpu.T = n
}

func (n *NSF) Play(samples int) []float32 {
	playDur := time.Duration(n.SpeedNTSC) * time.Nanosecond * 1000
	ticksPerPlay := int64(playDur / (time.Second / cpuClock))
	n.samples = make([]float32, 0, samples)
	for len(n.samples) < samples {
		n.playTicks = 0
		n.Cpu.PC = n.PlayAddr
		n.Cpu.Halt = false
		for !n.Cpu.Halt && len(n.samples) < samples {
			n.Cpu.Step()
			if !n.Cpu.I() {
				panic("INTERRUPT")
			}
		}
		for i := ticksPerPlay - n.playTicks; i > 0 && len(n.samples) < samples; i-- {
			n.Tick()
		}
	}
	return n.samples
}

// little-endian [2]byte to uint16 conversion
func bLEtoUint16(b []byte) uint16 {
	return uint16(b[1])<<8 + uint16(b[0])
}

// null-terminated bytes to string
func bToString(b []byte) string {
	i := 0
	for i = range b {
		if b[i] == 0 {
			break
		}
	}
	return string(b[:i])
}

type Ram struct {
	M [0xffff + 1]byte
	A Apu
}

func (r *Ram) Read(v uint16) byte {
	if v == 0x4015 {
		return r.A.Read(v)
	} else {
		return r.M[v]
	}
}

func (r *Ram) Write(v uint16, b byte) {
	r.M[v] = b
	if v&0xf000 == 0x4000 {
		r.A.Write(v, b)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2012-2018 Eli Janssen
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

// Package camo provides an HTTP proxy server with content type
// restrictions as well as regex host allow list support.
package camo

import (
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/cactus/go-camo/pkg/camo/encoding"
	"github.com/cactus/go-camo/pkg/htrie"

	"github.com/cactus/mlog"
)

// Config holds configuration data used when creating a Proxy with New.
type Config struct {
	// HMACKey is a byte slice to be used as the hmac key
	HMACKey []byte
	// Server name used in Headers and Via checks
	ServerName string
	// MaxSize is the maximum valid image size response (in bytes).
	MaxSize int64
	// MaxRedirects is the maximum number of redirects to follow.
	MaxRedirects int
	// Request timeout is a timeout for fetching upstream data.
	RequestTimeout time.Duration
	// Keepalive enable/disable
	DisableKeepAlivesFE bool
	DisableKeepAlivesBE bool
	// x-forwarded-for enable/disable
	EnableXFwdFor bool
	// additional content types to allow
	AllowContentVideo bool
	// allow URLs to contain user/pass credentials
	AllowCredetialURLs bool
	// no ip filtering (test mode)
	noIPFiltering bool
}

// The FilterFunc type is a function that validates a *url.URL
// A true value approves the url. A false value rejects the url.
type FilterFunc func(*url.URL) bool

// A Proxy is a Camo like HTTP proxy, that provides content type
// restrictions as well as regex host allow list support.
type Proxy struct {
	client            *http.Client
	config            *Config
	acceptTypesFilter *htrie.GlobPathChecker
	acceptTypesString string
	filters           []FilterFunc
	filtersLen        int
}

// ServerHTTP handles the client request, validates the request is validly
// HMAC signed, filters based on the Allow list, and then proxies
// valid requests to the desired endpoint. Responses are filtered for
// proper image content types.
func (p *Proxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	if p.config.DisableKeepAlivesFE {
		w.Header().Set("Connection", "close")
	}

	if req.Header.Get("Via") == p.config.ServerName {
		http.Error(w, "Request loop failure", http.StatusNotFound)
		return
	}

	// split path and get components
	components := strings.Split(req.URL.Path, "/")
	if len(components) < 3 {
		http.Error(w, "Malformed request path", http.StatusNotFound)
		return
	}
	sigHash, encodedURL := components[1], components[2]

	mlog.Debugm("client request", mlog.Map{"req": req})

	sURL, ok := encoding.DecodeURL(p.config.HMACKey, sigHash, encodedURL)
	if !ok {
		http.Error(w, "Bad Signature", http.StatusForbidden)
		return
	}

	mlog.Debugm("signed client url", mlog.Map{"url": sURL})

	u, err := url.Parse(sURL)
	if err != nil {
		mlog.Debugm("url parse error", mlog.Map{"err": err})
		http.Error(w, "Bad url", http.StatusBadRequest)
		return
	}

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

	nreq, err := http.NewRequest(req.Method, sURL, nil)
	if err != nil {
		mlog.Debugm("could not create NewRequest", mlog.Map{"err": err})
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	}

	// filter headers
	p.copyHeaders(&nreq.Header, &req.Header, &ValidReqHeaders)

	// x-forwarded-for (if appropriate)
	if p.config.EnableXFwdFor {
		xfwd4 := req.Header.Get("X-Forwarded-For")
		if xfwd4 == "" {
			hostIP, _, err := net.SplitHostPort(req.RemoteAddr)
			if err == nil {
				// add forwarded for header, as long as it isn't a private
				// ip address (use isRejectedIP to get private filtering for free)
				if ip := net.ParseIP(hostIP); ip != nil {
					if !isRejectedIP(ip) {
						nreq.Header.Add("X-Forwarded-For", hostIP)
					}
				}
			}
		} else {
			nreq.Header.Add("X-Forwarded-For", xfwd4)
		}
	}

	// add/squash an accept header if the client didn't send one
	nreq.Header.Set("Accept", p.acceptTypesString)

	nreq.Header.Add("User-Agent", p.config.ServerName)
	nreq.Header.Add("Via", p.config.ServerName)

	mlog.Debugm("built outgoing request", mlog.Map{"req": nreq})

	resp, err := p.client.Do(nreq)

	if resp != nil {
		defer resp.Body.Close()
	}

	if err != nil {
		mlog.Debugm("could not connect to endpoint", mlog.Map{"err": err})
		// this is a bit janky, but better than peeling off the
		// 3 layers of wrapped errors and trying to get to net.OpErr and
		// still having to rely on string comparison to find out if it is
		// a net.errClosing or not.
		errString := err.Error()
		// go 1.5 changes this to http.httpError
		// go 1.4 has this as net.OpError
		// and the error strings are different depending on which version too.
		if strings.Contains(errString, "timeout") || strings.Contains(errString, "Client.Timeout") {
			http.Error(w, "Error Fetching Resource", http.StatusGatewayTimeout)
		} else if strings.Contains(errString, "use of closed") {
			http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		} else if strings.Contains(errString, "BadRedirect:") {
			// Got a bad redirect
			mlog.Debugm("response from upstream", mlog.Map{"err": err})
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		} else {
			// some other error. call it a not found (camo compliant)
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		}
		return
	}

	mlog.Debugm("response from upstream", mlog.Map{"resp": resp})

	// check for too large a response
	if resp.ContentLength > p.config.MaxSize {
		mlog.Debugm("content length exceeded", mlog.Map{"url": sURL})
		http.Error(w, "Content length exceeded", http.StatusNotFound)
		return
	}

	switch resp.StatusCode {
	case 200, 206:
		contentType := resp.Header.Get("Content-Type")

		if contentType == "" {
			mlog.Debug("Empty content-type returned")
			http.Error(w, "Empty content-type returned", http.StatusBadRequest)
			return
		}

		if !p.acceptTypesFilter.CheckPath(contentType) {
			mlog.Debugm("Unsupported content-type returned", mlog.Map{"type": u})
			http.Error(w, "Unsupported content-type returned", http.StatusBadRequest)
			return
		}
	case 300:
		http.Error(w, "Multiple choices not supported", http.StatusNotFound)
		return
	case 301, 302, 303, 307:
		// if we get a redirect here, we either disabled following,
		// or followed until max depth and still got one (redirect loop)
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 304:
		h := w.Header()
		p.copyHeaders(&h, &resp.Header, &ValidRespHeaders)
		w.WriteHeader(304)
		return
	case 404:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 500, 502, 503, 504:
		// upstream errors should probably just 502. client can try later.
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	default:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	h := w.Header()
	p.copyHeaders(&h, &resp.Header, &ValidRespHeaders)
	w.WriteHeader(resp.StatusCode)

	// get a []byte from bufpool, and put it back on defer
	buf := *bufPool.Get().(*[]byte)
	defer bufPool.Put(&buf)

	// since this uses io.Copy/CopyBuffer from the respBody, it is streaming
	// from the request to the response. This means it will nearly
	// always end up with a chunked response.
	_, err = io.CopyBuffer(w, resp.Body, buf)
	if err != nil {
		// only log broken pipe errors at debug level
		if isBrokenPipe(err) {
			mlog.Debugm("error writing response", mlog.Map{"err": err})
		} else {
			// unknown error and not a broken pipe
			mlog.Printm("error writing response", mlog.Map{"err": err})
		}

		return
	}

	mlog.Debugm("response to client", mlog.Map{"resp": w})
}

func (p *Proxy) checkURL(reqURL *url.URL) error {
	// reject localhost urls
	// lower case for matching is done by CheckHostname, so no need to
	// ToLower here also
	uHostname := reqURL.Hostname()
	if uHostname == "" || localsFilter.CheckHostname(uHostname) {
		return errors.New("Bad url host")
	}

	// if not allowed, reject credentialed/userinfo urls
	if !p.config.AllowCredetialURLs && reqURL.User != nil {
		return errors.New("Userinfo URL rejected")
	}

	// ip/whitelist/blacklist filtering
	if !p.config.noIPFiltering {
		// filter out rejected networks
		if ip := net.ParseIP(uHostname); ip != nil {
			if isRejectedIP(ip) {
				return errors.New("Denylist host failure")
			}
		} else {
			if ips, err := net.LookupIP(uHostname); err == nil {
				for _, ip := range ips {
					if isRejectedIP(ip) {
						return errors.New("Denylist host failure")
					}
				}
			}
		}
	}

	// evaluate filters. first false value "fails"
	for i := 0; i < p.filtersLen; i++ {
		if !p.filters[i](reqURL) {
			return errors.New("Rejected due to filter-ruleset")
		}
	}

	return nil
}

// copy headers from src into dst
// empty filter map will result in no filtering being done
func (p *Proxy) copyHeaders(dst, src *http.Header, filter *map[string]bool) {
	f := *filter
	filtering := false
	if len(f) > 0 {
		filtering = true
	}

	for k, vv := range *src {
		if x, ok := f[k]; filtering && (!ok || !x) {
			continue
		}
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// NewWithFilters returns a new Proxy that utilises the passed in proxy filters.
// filters are evaluated in order, and the first false response from a filter
// function halts further evaluation and fails the request.
func NewWithFilters(pc Config, filters []FilterFunc) (*Proxy, error) {
	proxy, err := New(pc)
	if err != nil {
		return nil, err
	}

	filterFuncs := make([]FilterFunc, 0)
	// check for nil entries, and copy the slice in case the original
	// is mutated.
	for _, filter := range filters {
		if filter != nil {
			filterFuncs = append(filterFuncs, filter)
		}
	}
	proxy.filters = filterFuncs
	proxy.filtersLen = len(filterFuncs)
	return proxy, nil
}

// New returns a new Proxy. Returns an error if Proxy could not be constructed.
func New(pc Config) (*Proxy, error) {
	tr := &http.Transport{
		DialContext: (&net.Dialer{
			Timeout:   3 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		TLSHandshakeTimeout: 3 * time.Second,

		// max idle conns. Go DetaultTransport uses 100, which seems like a
		// fairly reasonable number. Very busy servers may wish to raise
		// or lower this value.
		MaxIdleConns:        100,
		MaxIdleConnsPerHost: 8,

		// more defaults from DefaultTransport, with a few tweaks
		IdleConnTimeout:       30 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,

		DisableKeepAlives: pc.DisableKeepAlivesBE,
		// no need for compression with images
		// some xml/svg can be compressed, but apparently some clients can
		// exhibit weird behavior when those are compressed
		DisableCompression: true,
	}

	client := &http.Client{
		Transport: tr,
		// timeout
		Timeout: pc.RequestTimeout,
	}

	acceptTypes := []string{"image/*"}
	// add additional accept types, if appropriate
	if pc.AllowContentVideo {
		acceptTypes = append(acceptTypes, "video/*")
	}

	// re-use the htrie glob path checker for accept types validation
	acceptTypesFilter := htrie.NewGlobPathChecker()
	for _, v := range acceptTypes {
		err := acceptTypesFilter.AddRule("|i|" + v)
		if err != nil {
			return nil, err
		}
	}

	p := &Proxy{
		client:            client,
		config:            &pc,
		acceptTypesString: strings.Join(acceptTypes, ", "),
		acceptTypesFilter: acceptTypesFilter,
	}

	client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		if len(via) >= pc.MaxRedirects {
			mlog.Debug("Got bad redirect: Too many redirects")
			return errors.New("BadRedirect: Too many redirects")
		}
		err := p.checkURL(req.URL)
		if err != nil {
			mlog.Debugm("Got bad redirect", mlog.Map{"url": req})
			return fmt.Errorf("BadRedirect: %s", err)
		}

		return nil
	}

	return p, nil

}
<commit_msg>fix space<commit_after>// Copyright (c) 2012-2018 Eli Janssen
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

// Package camo provides an HTTP proxy server with content type
// restrictions as well as regex host allow list support.
package camo

import (
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/cactus/go-camo/pkg/camo/encoding"
	"github.com/cactus/go-camo/pkg/htrie"

	"github.com/cactus/mlog"
)

// Config holds configuration data used when creating a Proxy with New.
type Config struct {
	// HMACKey is a byte slice to be used as the hmac key
	HMACKey []byte
	// Server name used in Headers and Via checks
	ServerName string
	// MaxSize is the maximum valid image size response (in bytes).
	MaxSize int64
	// MaxRedirects is the maximum number of redirects to follow.
	MaxRedirects int
	// Request timeout is a timeout for fetching upstream data.
	RequestTimeout time.Duration
	// Keepalive enable/disable
	DisableKeepAlivesFE bool
	DisableKeepAlivesBE bool
	// x-forwarded-for enable/disable
	EnableXFwdFor bool
	// additional content types to allow
	AllowContentVideo bool
	// allow URLs to contain user/pass credentials
	AllowCredetialURLs bool
	// no ip filtering (test mode)
	noIPFiltering bool
}

// The FilterFunc type is a function that validates a *url.URL
// A true value approves the url. A false value rejects the url.
type FilterFunc func(*url.URL) bool

// A Proxy is a Camo like HTTP proxy, that provides content type
// restrictions as well as regex host allow list support.
type Proxy struct {
	client            *http.Client
	config            *Config
	acceptTypesFilter *htrie.GlobPathChecker
	acceptTypesString string
	filters           []FilterFunc
	filtersLen        int
}

// ServerHTTP handles the client request, validates the request is validly
// HMAC signed, filters based on the Allow list, and then proxies
// valid requests to the desired endpoint. Responses are filtered for
// proper image content types.
func (p *Proxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	if p.config.DisableKeepAlivesFE {
		w.Header().Set("Connection", "close")
	}

	if req.Header.Get("Via") == p.config.ServerName {
		http.Error(w, "Request loop failure", http.StatusNotFound)
		return
	}

	// split path and get components
	components := strings.Split(req.URL.Path, "/")
	if len(components) < 3 {
		http.Error(w, "Malformed request path", http.StatusNotFound)
		return
	}
	sigHash, encodedURL := components[1], components[2]

	mlog.Debugm("client request", mlog.Map{"req": req})

	sURL, ok := encoding.DecodeURL(p.config.HMACKey, sigHash, encodedURL)
	if !ok {
		http.Error(w, "Bad Signature", http.StatusForbidden)
		return
	}

	mlog.Debugm("signed client url", mlog.Map{"url": sURL})

	u, err := url.Parse(sURL)
	if err != nil {
		mlog.Debugm("url parse error", mlog.Map{"err": err})
		http.Error(w, "Bad url", http.StatusBadRequest)
		return
	}

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

	nreq, err := http.NewRequest(req.Method, sURL, nil)
	if err != nil {
		mlog.Debugm("could not create NewRequest", mlog.Map{"err": err})
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	}

	// filter headers
	p.copyHeaders(&nreq.Header, &req.Header, &ValidReqHeaders)

	// x-forwarded-for (if appropriate)
	if p.config.EnableXFwdFor {
		xfwd4 := req.Header.Get("X-Forwarded-For")
		if xfwd4 == "" {
			hostIP, _, err := net.SplitHostPort(req.RemoteAddr)
			if err == nil {
				// add forwarded for header, as long as it isn't a private
				// ip address (use isRejectedIP to get private filtering for free)
				if ip := net.ParseIP(hostIP); ip != nil {
					if !isRejectedIP(ip) {
						nreq.Header.Add("X-Forwarded-For", hostIP)
					}
				}
			}
		} else {
			nreq.Header.Add("X-Forwarded-For", xfwd4)
		}
	}

	// add/squash an accept header if the client didn't send one
	nreq.Header.Set("Accept", p.acceptTypesString)

	nreq.Header.Add("User-Agent", p.config.ServerName)
	nreq.Header.Add("Via", p.config.ServerName)

	mlog.Debugm("built outgoing request", mlog.Map{"req": nreq})

	resp, err := p.client.Do(nreq)

	if resp != nil {
		defer resp.Body.Close()
	}

	if err != nil {
		mlog.Debugm("could not connect to endpoint", mlog.Map{"err": err})
		// this is a bit janky, but better than peeling off the
		// 3 layers of wrapped errors and trying to get to net.OpErr and
		// still having to rely on string comparison to find out if it is
		// a net.errClosing or not.
		errString := err.Error()
		// go 1.5 changes this to http.httpError
		// go 1.4 has this as net.OpError
		// and the error strings are different depending on which version too.
		if strings.Contains(errString, "timeout") || strings.Contains(errString, "Client.Timeout") {
			http.Error(w, "Error Fetching Resource", http.StatusGatewayTimeout)
		} else if strings.Contains(errString, "use of closed") {
			http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		} else if strings.Contains(errString, "BadRedirect:") {
			// Got a bad redirect
			mlog.Debugm("response from upstream", mlog.Map{"err": err})
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		} else {
			// some other error. call it a not found (camo compliant)
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		}
		return
	}

	mlog.Debugm("response from upstream", mlog.Map{"resp": resp})

	// check for too large a response
	if resp.ContentLength > p.config.MaxSize {
		mlog.Debugm("content length exceeded", mlog.Map{"url": sURL})
		http.Error(w, "Content length exceeded", http.StatusNotFound)
		return
	}

	switch resp.StatusCode {
	case 200, 206:
		contentType := resp.Header.Get("Content-Type")

		if contentType == "" {
			mlog.Debug("Empty content-type returned")
			http.Error(w, "Empty content-type returned", http.StatusBadRequest)
			return
		}

		if !p.acceptTypesFilter.CheckPath(contentType) {
			mlog.Debugm("Unsupported content-type returned", mlog.Map{"type": u})
			http.Error(w, "Unsupported content-type returned", http.StatusBadRequest)
			return
		}
	case 300:
		http.Error(w, "Multiple choices not supported", http.StatusNotFound)
		return
	case 301, 302, 303, 307:
		// if we get a redirect here, we either disabled following,
		// or followed until max depth and still got one (redirect loop)
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 304:
		h := w.Header()
		p.copyHeaders(&h, &resp.Header, &ValidRespHeaders)
		w.WriteHeader(304)
		return
	case 404:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 500, 502, 503, 504:
		// upstream errors should probably just 502. client can try later.
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	default:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	h := w.Header()
	p.copyHeaders(&h, &resp.Header, &ValidRespHeaders)
	w.WriteHeader(resp.StatusCode)

	// get a []byte from bufpool, and put it back on defer
	buf := *bufPool.Get().(*[]byte)
	defer bufPool.Put(&buf)

	// since this uses io.Copy/CopyBuffer from the respBody, it is streaming
	// from the request to the response. This means it will nearly
	// always end up with a chunked response.
	_, err = io.CopyBuffer(w, resp.Body, buf)
	if err != nil {
		// only log broken pipe errors at debug level
		if isBrokenPipe(err) {
			mlog.Debugm("error writing response", mlog.Map{"err": err})
		} else {
			// unknown error and not a broken pipe
			mlog.Printm("error writing response", mlog.Map{"err": err})
		}

		return
	}

	mlog.Debugm("response to client", mlog.Map{"resp": w})
}

func (p *Proxy) checkURL(reqURL *url.URL) error {
	// reject localhost urls
	// lower case for matching is done by CheckHostname, so no need to
	// ToLower here also
	uHostname := reqURL.Hostname()
	if uHostname == "" || localsFilter.CheckHostname(uHostname) {
		return errors.New("Bad url host")
	}

	// if not allowed, reject credentialed/userinfo urls
	if !p.config.AllowCredetialURLs && reqURL.User != nil {
		return errors.New("Userinfo URL rejected")
	}

	// ip/whitelist/blacklist filtering
	if !p.config.noIPFiltering {
		// filter out rejected networks
		if ip := net.ParseIP(uHostname); ip != nil {
			if isRejectedIP(ip) {
				return errors.New("Denylist host failure")
			}
		} else {
			if ips, err := net.LookupIP(uHostname); err == nil {
				for _, ip := range ips {
					if isRejectedIP(ip) {
						return errors.New("Denylist host failure")
					}
				}
			}
		}
	}

	// evaluate filters. first false value "fails"
	for i := 0; i < p.filtersLen; i++ {
		if !p.filters[i](reqURL) {
			return errors.New("Rejected due to filter-ruleset")
		}
	}

	return nil
}

// copy headers from src into dst
// empty filter map will result in no filtering being done
func (p *Proxy) copyHeaders(dst, src *http.Header, filter *map[string]bool) {
	f := *filter
	filtering := false
	if len(f) > 0 {
		filtering = true
	}

	for k, vv := range *src {
		if x, ok := f[k]; filtering && (!ok || !x) {
			continue
		}
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// NewWithFilters returns a new Proxy that utilises the passed in proxy filters.
// filters are evaluated in order, and the first false response from a filter
// function halts further evaluation and fails the request.
func NewWithFilters(pc Config, filters []FilterFunc) (*Proxy, error) {
	proxy, err := New(pc)
	if err != nil {
		return nil, err
	}

	filterFuncs := make([]FilterFunc, 0)
	// check for nil entries, and copy the slice in case the original
	// is mutated.
	for _, filter := range filters {
		if filter != nil {
			filterFuncs = append(filterFuncs, filter)
		}
	}
	proxy.filters = filterFuncs
	proxy.filtersLen = len(filterFuncs)
	return proxy, nil
}

// New returns a new Proxy. Returns an error if Proxy could not be constructed.
func New(pc Config) (*Proxy, error) {
	tr := &http.Transport{
		DialContext: (&net.Dialer{
			Timeout:   3 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		TLSHandshakeTimeout: 3 * time.Second,

		// max idle conns. Go DetaultTransport uses 100, which seems like a
		// fairly reasonable number. Very busy servers may wish to raise
		// or lower this value.
		MaxIdleConns:        100,
		MaxIdleConnsPerHost: 8,

		// more defaults from DefaultTransport, with a few tweaks
		IdleConnTimeout:       30 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,

		DisableKeepAlives: pc.DisableKeepAlivesBE,
		// no need for compression with images
		// some xml/svg can be compressed, but apparently some clients can
		// exhibit weird behavior when those are compressed
		DisableCompression: true,
	}

	client := &http.Client{
		Transport: tr,
		// timeout
		Timeout: pc.RequestTimeout,
	}

	acceptTypes := []string{"image/*"}
	// add additional accept types, if appropriate
	if pc.AllowContentVideo {
		acceptTypes = append(acceptTypes, "video/*")
	}

	// re-use the htrie glob path checker for accept types validation
	acceptTypesFilter := htrie.NewGlobPathChecker()
	for _, v := range acceptTypes {
		err := acceptTypesFilter.AddRule("|i|" + v)
		if err != nil {
			return nil, err
		}
	}

	p := &Proxy{
		client:            client,
		config:            &pc,
		acceptTypesString: strings.Join(acceptTypes, ", "),
		acceptTypesFilter: acceptTypesFilter,
	}

	client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		if len(via) >= pc.MaxRedirects {
			mlog.Debug("Got bad redirect: Too many redirects")
			return errors.New("BadRedirect: Too many redirects")
		}
		err := p.checkURL(req.URL)
		if err != nil {
			mlog.Debugm("Got bad redirect", mlog.Map{"url": req})
			return fmt.Errorf("BadRedirect: %s", err)
		}

		return nil
	}

	return p, nil
}
<|endoftext|>


<commit_before>package cli

import (
	"errors"
	"fmt"
	"io"
	"math/rand"
	"strconv"

	"github.com/apprenda/kismatic/pkg/install"
	"github.com/spf13/cobra"
)

// NewCmdVolume returns the storage command
func NewCmdVolume(out io.Writer) *cobra.Command {
	var planFile string
	cmd := &cobra.Command{
		Use:   "volume",
		Short: "manage storage volumes on your Kubernetes cluster",
		RunE: func(cmd *cobra.Command, args []string) error {
			return cmd.Usage()
		},
	}
	addPlanFileFlag(cmd.PersistentFlags(), &planFile)
	cmd.AddCommand(NewCmdVolumeAdd(out, &planFile))
	return cmd
}

type volumeAddOptions struct {
	replicaCount       int
	distributionCount  int
	allowAddress       []string
	verbose            bool
	outputFormat       string
	generatedAssetsDir string
}

// NewCmdVolumeAdd returns the command for adding storage volumes
func NewCmdVolumeAdd(out io.Writer, planFile *string) *cobra.Command {
	opts := volumeAddOptions{}
	cmd := &cobra.Command{
		Use:   "add size_in_gigabytes [volume name]",
		Short: "add storage volumes to the Kubernetes cluster",
		Long: `Add storage volumes to the Kubernetes cluster.

This function requires a target cluster that has storage nodes.`,
		RunE: func(cmd *cobra.Command, args []string) error {
			return doVolumeAdd(out, opts, *planFile, args)
		},
		Example: `  Create a distributed, replicated volume,
  named "storage01" with a 10 GB quota,
  grant access to any IP address starting with 10.10.
  kismatic volume add 10 storage01 -r 2 -d 2 -a 10.10.*.*
		`,
	}
	cmd.Flags().IntVarP(&opts.replicaCount, "replica-count", "r", 2, "The number of times each file will be written.")
	cmd.Flags().IntVarP(&opts.distributionCount, "distribution-count", "d", 1, "This is the degree to which data will be distributed across the cluster. By default, it won't be -- each replica will receive 100% of the data. Distribution makes listing or backing up the cluster more complicated by spreading data around the cluster but makes reads and writes more performant.")
	cmd.Flags().StringSliceVarP(&opts.allowAddress, "allow-address", "a", nil, "Comma delimited list of address wildcards permitted access to the volume in addition to Kubernetes nodes.")
	cmd.Flags().BoolVar(&opts.verbose, "verbose", false, "enable verbose logging")
	cmd.Flags().StringVarP(&opts.outputFormat, "output", "o", "simple", `output format (options "simple"|"raw")`)
	cmd.Flags().StringVar(&opts.generatedAssetsDir, "generated-assets-dir", "generated", "path to the directory where assets generated during the installation process will be stored")
	return cmd
}

func doVolumeAdd(out io.Writer, opts volumeAddOptions, planFile string, args []string) error {
	var volumeName string
	var volumeSizeStrGB string
	switch len(args) {
	case 0:
		return errors.New("the volume size (in gigabytes) must be provided as the first argument to add")
	case 1:
		volumeSizeStrGB = args[0]
		volumeName = "kismatic-" + generateRandomString(5)
	case 2:
		volumeSizeStrGB = args[0]
		volumeName = args[1]
	default:
		return fmt.Errorf("%d arguments were provided, but add does not support more than two arguments", len(args))
	}
	volumeSizeGB, err := strconv.Atoi(volumeSizeStrGB)
	if err != nil {
		return errors.New("the volume size provided is not valid")
	}
	// Setup ansible
	planner := &install.FilePlanner{File: planFile}
	if !planner.PlanExists() {
		return fmt.Errorf("Plan file not found at %q", planFile)
	}
	execOpts := install.ExecutorOptions{
		OutputFormat: opts.outputFormat,
		Verbose:      opts.verbose,
		// Need to refactor executor code... this will do for now as we don't need the generated assets dir in this command
		GeneratedAssetsDirectory: opts.generatedAssetsDir,
	}
	exec, err := install.NewExecutor(out, out, execOpts)
	if err != nil {
		return err
	}
	plan, err := planner.Read()
	if err != nil {
		return err
	}

	// Run validation
	vopts := &validateOpts{
<<<<<<< HEAD
		outputFormat:       opts.outputFormat,
		verbose:            opts.verbose,
		planFile:           planFile,
		skipPreFlight:      true,
		generatedAssetsDir: opts.generatedAssetsDir,
=======
		outputFormat:  opts.outputFormat,
		verbose:       opts.verbose,
		planFile:      planFile,
		skipPreFlight: true,
>>>>>>> abff976bcc8e149fc2c7302a96b3ba3cd6dcf0e9
	}
	if err := doValidate(out, planner, vopts); err != nil {
		return err
	}

	v := install.StorageVolume{
		Name:              volumeName,
		SizeGB:            volumeSizeGB,
		ReplicateCount:    opts.replicaCount,
		DistributionCount: opts.distributionCount,
	}
	if opts.allowAddress != nil {
		v.AllowAddresses = opts.allowAddress
	}
	if ok, errs := install.ValidateStorageVolume(v); !ok {
		fmt.Println("The storage volume configuration is not valid:")
		for _, e := range errs {
			fmt.Printf("- %s\n", e)
		}
		return errors.New("storage volume validation failed")
	}
	if err := exec.AddVolume(plan, v); err != nil {
		return fmt.Errorf("error adding new volume: %v", err)
	}

	fmt.Fprintln(out)
	fmt.Fprintln(out, "Successfully added the persistent volume to the kubernetes cluster.")
	fmt.Fprintln(out)
	fmt.Fprintf(out, "Use \"kubectl describe pv %s\" to view volume details.\n", v.Name)
	return nil
}

func generateRandomString(n int) string {
	// removed 1, l, o, 0 to prevent confusion
	chars := []rune("abcdefghijkmnpqrstuvwxyz23456789")
	res := make([]rune, n)
	for i := range res {
		res[i] = chars[rand.Intn(len(chars))]
	}
	return string(res)
}
<commit_msg>Remove merge conflict text<commit_after>package cli

import (
	"errors"
	"fmt"
	"io"
	"math/rand"
	"strconv"

	"github.com/apprenda/kismatic/pkg/install"
	"github.com/spf13/cobra"
)

// NewCmdVolume returns the storage command
func NewCmdVolume(out io.Writer) *cobra.Command {
	var planFile string
	cmd := &cobra.Command{
		Use:   "volume",
		Short: "manage storage volumes on your Kubernetes cluster",
		RunE: func(cmd *cobra.Command, args []string) error {
			return cmd.Usage()
		},
	}
	addPlanFileFlag(cmd.PersistentFlags(), &planFile)
	cmd.AddCommand(NewCmdVolumeAdd(out, &planFile))
	return cmd
}

type volumeAddOptions struct {
	replicaCount       int
	distributionCount  int
	allowAddress       []string
	verbose            bool
	outputFormat       string
	generatedAssetsDir string
}

// NewCmdVolumeAdd returns the command for adding storage volumes
func NewCmdVolumeAdd(out io.Writer, planFile *string) *cobra.Command {
	opts := volumeAddOptions{}
	cmd := &cobra.Command{
		Use:   "add size_in_gigabytes [volume name]",
		Short: "add storage volumes to the Kubernetes cluster",
		Long: `Add storage volumes to the Kubernetes cluster.

This function requires a target cluster that has storage nodes.`,
		RunE: func(cmd *cobra.Command, args []string) error {
			return doVolumeAdd(out, opts, *planFile, args)
		},
		Example: `  Create a distributed, replicated volume,
  named "storage01" with a 10 GB quota,
  grant access to any IP address starting with 10.10.
  kismatic volume add 10 storage01 -r 2 -d 2 -a 10.10.*.*
		`,
	}
	cmd.Flags().IntVarP(&opts.replicaCount, "replica-count", "r", 2, "The number of times each file will be written.")
	cmd.Flags().IntVarP(&opts.distributionCount, "distribution-count", "d", 1, "This is the degree to which data will be distributed across the cluster. By default, it won't be -- each replica will receive 100% of the data. Distribution makes listing or backing up the cluster more complicated by spreading data around the cluster but makes reads and writes more performant.")
	cmd.Flags().StringSliceVarP(&opts.allowAddress, "allow-address", "a", nil, "Comma delimited list of address wildcards permitted access to the volume in addition to Kubernetes nodes.")
	cmd.Flags().BoolVar(&opts.verbose, "verbose", false, "enable verbose logging")
	cmd.Flags().StringVarP(&opts.outputFormat, "output", "o", "simple", `output format (options "simple"|"raw")`)
	cmd.Flags().StringVar(&opts.generatedAssetsDir, "generated-assets-dir", "generated", "path to the directory where assets generated during the installation process will be stored")
	return cmd
}

func doVolumeAdd(out io.Writer, opts volumeAddOptions, planFile string, args []string) error {
	var volumeName string
	var volumeSizeStrGB string
	switch len(args) {
	case 0:
		return errors.New("the volume size (in gigabytes) must be provided as the first argument to add")
	case 1:
		volumeSizeStrGB = args[0]
		volumeName = "kismatic-" + generateRandomString(5)
	case 2:
		volumeSizeStrGB = args[0]
		volumeName = args[1]
	default:
		return fmt.Errorf("%d arguments were provided, but add does not support more than two arguments", len(args))
	}
	volumeSizeGB, err := strconv.Atoi(volumeSizeStrGB)
	if err != nil {
		return errors.New("the volume size provided is not valid")
	}
	// Setup ansible
	planner := &install.FilePlanner{File: planFile}
	if !planner.PlanExists() {
		return fmt.Errorf("Plan file not found at %q", planFile)
	}
	execOpts := install.ExecutorOptions{
		OutputFormat: opts.outputFormat,
		Verbose:      opts.verbose,
		// Need to refactor executor code... this will do for now as we don't need the generated assets dir in this command
		GeneratedAssetsDirectory: opts.generatedAssetsDir,
	}
	exec, err := install.NewExecutor(out, out, execOpts)
	if err != nil {
		return err
	}
	plan, err := planner.Read()
	if err != nil {
		return err
	}

	// Run validation
	vopts := &validateOpts{
		outputFormat:       opts.outputFormat,
		verbose:            opts.verbose,
		planFile:           planFile,
		skipPreFlight:      true,
		generatedAssetsDir: opts.generatedAssetsDir,
	}
	if err := doValidate(out, planner, vopts); err != nil {
		return err
	}

	v := install.StorageVolume{
		Name:              volumeName,
		SizeGB:            volumeSizeGB,
		ReplicateCount:    opts.replicaCount,
		DistributionCount: opts.distributionCount,
	}
	if opts.allowAddress != nil {
		v.AllowAddresses = opts.allowAddress
	}
	if ok, errs := install.ValidateStorageVolume(v); !ok {
		fmt.Println("The storage volume configuration is not valid:")
		for _, e := range errs {
			fmt.Printf("- %s\n", e)
		}
		return errors.New("storage volume validation failed")
	}
	if err := exec.AddVolume(plan, v); err != nil {
		return fmt.Errorf("error adding new volume: %v", err)
	}

	fmt.Fprintln(out)
	fmt.Fprintln(out, "Successfully added the persistent volume to the kubernetes cluster.")
	fmt.Fprintln(out)
	fmt.Fprintf(out, "Use \"kubectl describe pv %s\" to view volume details.\n", v.Name)
	return nil
}

func generateRandomString(n int) string {
	// removed 1, l, o, 0 to prevent confusion
	chars := []rune("abcdefghijkmnpqrstuvwxyz23456789")
	res := make([]rune, n)
	for i := range res {
		res[i] = chars[rand.Intn(len(chars))]
	}
	return string(res)
}
<|endoftext|>


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

package log

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"sort"
	"strings"

	"github.com/sirupsen/logrus"
	"gopkg.in/natefinch/lumberjack.v2"
)

// Formatter is for user to specific the log formatter
type Formatter int

const (
	// TextFormatter means unordered fields
	TextFormatter Formatter = iota
	// OrderedTextFormatter writes the fields(only but not level&msg) orderly
	OrderedTextFormatter
	// JSON, easy to support but we don't need it right now
)

// InitLogger set the output and formatter
func InitLogger(filename string, f Formatter) {
	setOutput(filename)
	if f == OrderedTextFormatter {
		fm := &orderedFieldsTextFormatter{}
		logrus.SetFormatter(fm)
	}
}

// setOutput sets log output to filename globally.
// filename="/var/log/sqlflow.log": write the log to file
// filename="": write the file to stdout or stderr
// filename="/dev/null": ignore log message
func setOutput(filename string) {
	filename = strings.Trim(filename, " ")
	if filename == "/dev/null" {
		logrus.SetOutput(ioutil.Discard)
	} else if len(filename) > 0 {
		logrus.SetOutput(&lumberjack.Logger{
			Filename:   filename,
			MaxSize:    32, // megabytes
			MaxBackups: 64,
			MaxAge:     15, // days
			Compress:   true,
		})
	}
}

// orderedFieldsTextFormatter writes the fields(only but not level or msg) orderly
type orderedFieldsTextFormatter struct {
}

func (f *orderedFieldsTextFormatter) Format(logger *logrus.Entry) ([]byte, error) {
	var b *bytes.Buffer
	if logger.Buffer != nil {
		b = logger.Buffer
	} else {
		b = &bytes.Buffer{}
	}
	fmt.Fprintf(b, "%s %s msg=\"%s\"", logger.Time.Format("2006-01-02 15:04:05"), logger.Level.String(), logger.Message)

	keys := make([]string, 0, len(logger.Data))
	for k := range logger.Data {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		v := logger.Data[k]
		_, ok := v.(string)
		if ok {
			fmt.Fprintf(b, " %s=\"%s\"", k, v)
		} else {
			fmt.Fprintf(b, " %s=%v", k, v)
		}
	}
	b.WriteByte('\n')
	return b.Bytes(), nil
}
<commit_msg>Enable log in sqlflowserver (#2550)<commit_after>// Copyright 2020 The SQLFlow Authors. All rights reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package log

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"sort"
	"strings"

	"github.com/sirupsen/logrus"
	"gopkg.in/natefinch/lumberjack.v2"
)

// Formatter is for user to specific the log formatter
type Formatter int

const (
	// TextFormatter means unordered fields
	TextFormatter Formatter = iota
	// OrderedTextFormatter writes the fields(only but not level&msg) orderly
	OrderedTextFormatter
	// JSON, easy to support but we don't need it right now
)

// InitLogger set the output and formatter
func InitLogger(filename string, f Formatter) {
	setOutput(filename)
	if f == OrderedTextFormatter {
		fm := &orderedFieldsTextFormatter{}
		logrus.SetFormatter(fm)
	}
}

// setOutput sets log output to filename globally.
// filename="/var/log/sqlflow.log": write the log to file
// filename="": write the file to stdout or stderr
// filename="/dev/null": ignore log message
func setOutput(filename string) {
	filename = strings.Trim(filename, " ")
	if filename == "/dev/null" {
		logrus.SetOutput(ioutil.Discard)
	} else if filename == "" {
		logrus.SetOutput(os.Stdout)
	} else if len(filename) > 0 {
		logrus.SetOutput(&lumberjack.Logger{
			Filename:   filename,
			MaxSize:    32, // megabytes
			MaxBackups: 64,
			MaxAge:     15, // days
			Compress:   true,
		})
	}
}

// orderedFieldsTextFormatter writes the fields(only but not level or msg) orderly
type orderedFieldsTextFormatter struct {
}

func (f *orderedFieldsTextFormatter) Format(logger *logrus.Entry) ([]byte, error) {
	var b *bytes.Buffer
	if logger.Buffer != nil {
		b = logger.Buffer
	} else {
		b = &bytes.Buffer{}
	}
	fmt.Fprintf(b, "%s %s msg=\"%s\"", logger.Time.Format("2006-01-02 15:04:05"), logger.Level.String(), logger.Message)

	keys := make([]string, 0, len(logger.Data))
	for k := range logger.Data {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		v := logger.Data[k]
		_, ok := v.(string)
		if ok {
			fmt.Fprintf(b, " %s=\"%s\"", k, v)
		} else {
			fmt.Fprintf(b, " %s=%v", k, v)
		}
	}
	b.WriteByte('\n')
	return b.Bytes(), nil
}
<|endoftext|>