<commit_before>/*
 * Minio Client, (C) 2015, 2016 Minio, 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 cmd

import (
	"encoding/json"
	"strings"

	"github.com/fatih/color"
	"github.com/minio/cli"
	"github.com/minio/mc/pkg/console"
	"github.com/minio/minio/pkg/probe"
)

var (
	policyFlags = []cli.Flag{
		cli.BoolFlag{
			Name:  "recursive, r",
			Usage: "List recursively.",
		},
	}
)

// Set public policy
var policyCmd = cli.Command{
	Name:   "policy",
	Usage:  "Manage anonymous access to objects.",
	Action: mainPolicy,
	Flags:  append(policyFlags, globalFlags...),
	CustomHelpTemplate: `Name:
   mc {{.Name}} - {{.Usage}}

USAGE:
   mc {{.Name}} [FLAGS] PERMISSION TARGET
   mc {{.Name}} [FLAGS] TARGET

PERMISSION:
   Allowed policies are: [none, download, upload, public].

FLAGS:
  {{range .Flags}}{{.}}
  {{end}}
EXAMPLES:
   1. Set bucket to "download" on Amazon S3 cloud storage.
      $ mc {{.Name}} download s3/burningman2011

   2. Set bucket to "public" on Amazon S3 cloud storage.
      $ mc {{.Name}} public s3/shared

   3. Set bucket to "upload" on Amazon S3 cloud storage.
      $ mc {{.Name}} upload s3/incoming

   4. Set a prefix to "public" on Amazon S3 cloud storage.
      $ mc {{.Name}} public s3/public-commons/images

   5. Get bucket permissions.
      $ mc {{.Name}} s3/shared

   6. List policies set to a specified bucket.
      $ mc {{.Name}} list s3/shared

   7. List public object URLs recursively.
      $ mc {{.Name}} --recursive links s3/shared/

`,
}

// policyRules contains policy rule
type policyRules struct {
	Resource string `json:"resource"`
	Allow    string `json:"allow"`
}

// String colorized access message.
func (s policyRules) String() string {
	return console.Colorize("Policy", s.Resource+" => "+s.Allow+"")
}

// JSON jsonified policy message.
func (s policyRules) JSON() string {
	policyJSONBytes, e := json.Marshal(s)
	fatalIf(probe.NewError(e), "Unable to marshal into JSON.")
	return string(policyJSONBytes)
}

// policyMessage is container for policy command on bucket success and failure messages.
type policyMessage struct {
	Operation string      `json:"operation"`
	Status    string      `json:"status"`
	Bucket    string      `json:"bucket"`
	Perms     accessPerms `json:"permission"`
}

// String colorized access message.
func (s policyMessage) String() string {
	if s.Operation == "set" {
		return console.Colorize("Policy",
			"Access permission for ‘"+s.Bucket+"’ is set to ‘"+string(s.Perms)+"’")
	}
	if s.Operation == "get" {
		return console.Colorize("Policy",
			"Access permission for ‘"+s.Bucket+"’"+" is ‘"+string(s.Perms)+"’")
	}
	// nothing to print
	return ""
}

// JSON jsonified policy message.
func (s policyMessage) JSON() string {
	policyJSONBytes, e := json.Marshal(s)
	fatalIf(probe.NewError(e), "Unable to marshal into JSON.")

	return string(policyJSONBytes)
}

// policyLinksMessage is container for policy links command
type policyLinksMessage struct {
	Status string `json:"status"`
	URL    string `json:"url"`
}

// String colorized access message.
func (s policyLinksMessage) String() string {
	return console.Colorize("Policy", string(s.URL))
}

// JSON jsonified policy message.
func (s policyLinksMessage) JSON() string {
	policyJSONBytes, e := json.Marshal(s)
	fatalIf(probe.NewError(e), "Unable to marshal into JSON.")

	return string(policyJSONBytes)
}

// checkPolicySyntax check for incoming syntax.
func checkPolicySyntax(ctx *cli.Context) {
	argsLength := len(ctx.Args())
	// Always print a help message when we have extra arguments
	if argsLength > 2 {
		cli.ShowCommandHelpAndExit(ctx, "policy", 1) // last argument is exit code.
	}
	// Always print a help message when no arguments specified
	if argsLength < 1 {
		cli.ShowCommandHelpAndExit(ctx, "policy", 1)
	}

	firstArg := ctx.Args().Get(0)

	// More syntax checking
	switch accessPerms(firstArg) {
	case accessNone, accessDownload, accessUpload, accessPublic:
		// Always expect two arguments when a policy permission is provided
		if argsLength != 2 {
			cli.ShowCommandHelpAndExit(ctx, "policy", 1)
		}
	case "list":
		// Always expect an argument after list cmd
		if argsLength != 2 {
			cli.ShowCommandHelpAndExit(ctx, "policy", 1)
		}
	case "links":
		// Always expect an argument after links cmd
		if argsLength != 2 {
			cli.ShowCommandHelpAndExit(ctx, "policy", 1)
		}

	default:
		if argsLength == 2 {
			fatalIf(errDummy().Trace(),
				"Unrecognized permission ‘"+string(firstArg)+"’. Allowed values are [none, download, upload, public].")
		}
	}
}

// Convert an accessPerms to a string recognizable by minio-go
func accessPermToString(perm accessPerms) string {
	policy := ""
	switch perm {
	case accessNone:
		policy = "none"
	case accessDownload:
		policy = "readonly"
	case accessUpload:
		policy = "writeonly"
	case accessPublic:
		policy = "readwrite"
	}
	return policy
}

// doSetAccess do set access.
func doSetAccess(targetURL string, targetPERMS accessPerms) *probe.Error {
	clnt, err := newClient(targetURL)
	if err != nil {
		return err.Trace(targetURL)
	}
	policy := accessPermToString(targetPERMS)
	if err = clnt.SetAccess(policy); err != nil {
		return err.Trace(targetURL, string(targetPERMS))
	}
	return nil
}

// Convert a minio-go permission to accessPerms type
func stringToAccessPerm(perm string) accessPerms {
	var policy accessPerms
	switch perm {
	case "none":
		policy = accessNone
	case "readonly":
		policy = accessDownload
	case "writeonly":
		policy = accessUpload
	case "readwrite":
		policy = accessPublic
	}
	return policy
}

// doGetAccess do get access.
func doGetAccess(targetURL string) (perms accessPerms, err *probe.Error) {
	clnt, err := newClient(targetURL)
	if err != nil {
		return "", err.Trace(targetURL)
	}
	perm, err := clnt.GetAccess()
	if err != nil {
		return "", err.Trace(targetURL)
	}
	return stringToAccessPerm(perm), nil
}

// doGetAccessRules do get access rules.
func doGetAccessRules(targetURL string) (r map[string]string, err *probe.Error) {
	clnt, err := newClient(targetURL)
	if err != nil {
		return map[string]string{}, err.Trace(targetURL)
	}
	return clnt.GetAccessRules()
}

// Run policy list command
func runPolicyListCmd(ctx *cli.Context) {
	targetURL := ctx.Args().Last()
	policies, err := doGetAccessRules(targetURL)
	if err != nil {
		switch err.ToGoError().(type) {
		case APINotImplemented:
			fatalIf(err.Trace(), "Unable to list policies of a non S3 url ‘"+targetURL+"’.")
		default:
			fatalIf(err.Trace(targetURL), "Unable to list policies of target ‘"+targetURL+"’.")
		}
	}
	for k, v := range policies {
		printMsg(policyRules{Resource: k, Allow: v})
	}
}

// Run policy links command
func runPolicyLinksCmd(ctx *cli.Context) {
	// Get alias/bucket/prefix argument
	targetURL := ctx.Args().Last()

	// Fetch all policies associated to the passed url
	policies, err := doGetAccessRules(targetURL)
	if err != nil {
		switch err.ToGoError().(type) {
		case APINotImplemented:
			fatalIf(err.Trace(), "Unable to list policies of a non S3 url ‘"+targetURL+"’.")
		default:
			fatalIf(err.Trace(targetURL), "Unable to list policies of target ‘"+targetURL+"’.")
		}
	}

	// Extract alias from the passed argument, we'll need it to
	// construct new pathes to list public objects
	alias, path := url2Alias(targetURL)

	isRecursive := ctx.Bool("recursive")
	isIncomplete := false

	// Iterate over policy rules to fetch public urls, then search
	// for objects under those urls
	for k, v := range policies {
		// Trim the asterisk in policy rules
		policyPath := strings.TrimSuffix(k, "*")
		// Check if current policy prefix is related to the url passed by the user
		if !strings.HasPrefix(policyPath, path) {
			continue
		}
		// Check if the found policy has read permission
		perm := stringToAccessPerm(v)
		if perm != accessDownload && perm != accessPublic {
			continue
		}
		// Construct the new path to search for public objects
		newURL := alias + "/" + policyPath
		clnt, err := newClient(newURL)
		fatalIf(err.Trace(newURL), "Unable to initialize target ‘"+targetURL+"’.")
		// Search for public objects
		for content := range clnt.List(isRecursive, isIncomplete, DirFirst) {
			if content.Err != nil {
				errorIf(content.Err.Trace(clnt.GetURL().String()), "Unable to list folder.")
				continue
			}
			// Construct the message to be displayed to the user
			msg := policyLinksMessage{
				Status: "success",
				URL:    content.URL.String(),
			}
			// Print the found object
			printMsg(msg)
		}
	}
}

// Run policy cmd to fetch set permission
func runPolicyCmd(ctx *cli.Context) {
	perms := accessPerms(ctx.Args().First())
	if perms.isValidAccessPERM() {
		targetURL := ctx.Args().Last()
		err := doSetAccess(targetURL, perms)
		// Upon error exit.
		if err != nil {
			switch err.ToGoError().(type) {
			case APINotImplemented:
				fatalIf(err.Trace(), "Unable to set policy of a non S3 url ‘"+targetURL+"’.")
			default:
				fatalIf(err.Trace(targetURL, string(perms)),
					"Unable to set policy ‘"+string(perms)+"’ for ‘"+targetURL+"’.")

			}
		}

		printMsg(policyMessage{
			Status:    "success",
			Operation: "set",
			Bucket:    targetURL,
			Perms:     perms,
		})
	} else {
		targetURL := ctx.Args().First()
		perms, err := doGetAccess(targetURL)
		if err != nil {
			switch err.ToGoError().(type) {
			case APINotImplemented:
				fatalIf(err.Trace(), "Unable to get policy of a non S3 url ‘"+targetURL+"’.")
			default:
				fatalIf(err.Trace(targetURL), "Unable to get policy for ‘"+targetURL+"’.")
			}
		}

		printMsg(policyMessage{
			Status:    "success",
			Operation: "get",
			Bucket:    targetURL,
			Perms:     perms,
		})
	}
}

func mainPolicy(ctx *cli.Context) error {
	// Set global flags from context.
	setGlobalsFromContext(ctx)

	// check 'policy' cli arguments.
	checkPolicySyntax(ctx)

	// Additional command speific theme customization.
	console.SetColor("Policy", color.New(color.FgGreen, color.Bold))

	switch ctx.Args().First() {
	case "list":
		// policy list alias/bucket/prefix
		runPolicyListCmd(ctx)
	case "links":
		// policy links alias/bucket/prefix
		runPolicyLinksCmd(ctx)
	default:
		// policy [download|upload|public|] alias/bucket/prefix
		runPolicyCmd(ctx)
	}
	return nil
}
<commit_msg>links: Print encoded form of public URLs (#1947)<commit_after>/*
 * Minio Client, (C) 2015, 2016 Minio, 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 cmd

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

	"github.com/fatih/color"
	"github.com/minio/cli"
	"github.com/minio/mc/pkg/console"
	"github.com/minio/minio/pkg/probe"
)

var (
	policyFlags = []cli.Flag{
		cli.BoolFlag{
			Name:  "recursive, r",
			Usage: "List recursively.",
		},
	}
)

// Set public policy
var policyCmd = cli.Command{
	Name:   "policy",
	Usage:  "Manage anonymous access to objects.",
	Action: mainPolicy,
	Flags:  append(policyFlags, globalFlags...),
	CustomHelpTemplate: `Name:
   mc {{.Name}} - {{.Usage}}

USAGE:
   mc {{.Name}} [FLAGS] PERMISSION TARGET
   mc {{.Name}} [FLAGS] TARGET

PERMISSION:
   Allowed policies are: [none, download, upload, public].

FLAGS:
  {{range .Flags}}{{.}}
  {{end}}
EXAMPLES:
   1. Set bucket to "download" on Amazon S3 cloud storage.
      $ mc {{.Name}} download s3/burningman2011

   2. Set bucket to "public" on Amazon S3 cloud storage.
      $ mc {{.Name}} public s3/shared

   3. Set bucket to "upload" on Amazon S3 cloud storage.
      $ mc {{.Name}} upload s3/incoming

   4. Set a prefix to "public" on Amazon S3 cloud storage.
      $ mc {{.Name}} public s3/public-commons/images

   5. Get bucket permissions.
      $ mc {{.Name}} s3/shared

   6. List policies set to a specified bucket.
      $ mc {{.Name}} list s3/shared

   7. List public object URLs recursively.
      $ mc {{.Name}} --recursive links s3/shared/

`,
}

// policyRules contains policy rule
type policyRules struct {
	Resource string `json:"resource"`
	Allow    string `json:"allow"`
}

// String colorized access message.
func (s policyRules) String() string {
	return console.Colorize("Policy", s.Resource+" => "+s.Allow+"")
}

// JSON jsonified policy message.
func (s policyRules) JSON() string {
	policyJSONBytes, e := json.Marshal(s)
	fatalIf(probe.NewError(e), "Unable to marshal into JSON.")
	return string(policyJSONBytes)
}

// policyMessage is container for policy command on bucket success and failure messages.
type policyMessage struct {
	Operation string      `json:"operation"`
	Status    string      `json:"status"`
	Bucket    string      `json:"bucket"`
	Perms     accessPerms `json:"permission"`
}

// String colorized access message.
func (s policyMessage) String() string {
	if s.Operation == "set" {
		return console.Colorize("Policy",
			"Access permission for ‘"+s.Bucket+"’ is set to ‘"+string(s.Perms)+"’")
	}
	if s.Operation == "get" {
		return console.Colorize("Policy",
			"Access permission for ‘"+s.Bucket+"’"+" is ‘"+string(s.Perms)+"’")
	}
	// nothing to print
	return ""
}

// JSON jsonified policy message.
func (s policyMessage) JSON() string {
	policyJSONBytes, e := json.Marshal(s)
	fatalIf(probe.NewError(e), "Unable to marshal into JSON.")

	return string(policyJSONBytes)
}

// policyLinksMessage is container for policy links command
type policyLinksMessage struct {
	Status string `json:"status"`
	URL    string `json:"url"`
}

// String colorized access message.
func (s policyLinksMessage) String() string {
	return console.Colorize("Policy", string(s.URL))
}

// JSON jsonified policy message.
func (s policyLinksMessage) JSON() string {
	policyJSONBytes, e := json.Marshal(s)
	fatalIf(probe.NewError(e), "Unable to marshal into JSON.")

	return string(policyJSONBytes)
}

// checkPolicySyntax check for incoming syntax.
func checkPolicySyntax(ctx *cli.Context) {
	argsLength := len(ctx.Args())
	// Always print a help message when we have extra arguments
	if argsLength > 2 {
		cli.ShowCommandHelpAndExit(ctx, "policy", 1) // last argument is exit code.
	}
	// Always print a help message when no arguments specified
	if argsLength < 1 {
		cli.ShowCommandHelpAndExit(ctx, "policy", 1)
	}

	firstArg := ctx.Args().Get(0)

	// More syntax checking
	switch accessPerms(firstArg) {
	case accessNone, accessDownload, accessUpload, accessPublic:
		// Always expect two arguments when a policy permission is provided
		if argsLength != 2 {
			cli.ShowCommandHelpAndExit(ctx, "policy", 1)
		}
	case "list":
		// Always expect an argument after list cmd
		if argsLength != 2 {
			cli.ShowCommandHelpAndExit(ctx, "policy", 1)
		}
	case "links":
		// Always expect an argument after links cmd
		if argsLength != 2 {
			cli.ShowCommandHelpAndExit(ctx, "policy", 1)
		}

	default:
		if argsLength == 2 {
			fatalIf(errDummy().Trace(),
				"Unrecognized permission ‘"+string(firstArg)+"’. Allowed values are [none, download, upload, public].")
		}
	}
}

// Convert an accessPerms to a string recognizable by minio-go
func accessPermToString(perm accessPerms) string {
	policy := ""
	switch perm {
	case accessNone:
		policy = "none"
	case accessDownload:
		policy = "readonly"
	case accessUpload:
		policy = "writeonly"
	case accessPublic:
		policy = "readwrite"
	}
	return policy
}

// doSetAccess do set access.
func doSetAccess(targetURL string, targetPERMS accessPerms) *probe.Error {
	clnt, err := newClient(targetURL)
	if err != nil {
		return err.Trace(targetURL)
	}
	policy := accessPermToString(targetPERMS)
	if err = clnt.SetAccess(policy); err != nil {
		return err.Trace(targetURL, string(targetPERMS))
	}
	return nil
}

// Convert a minio-go permission to accessPerms type
func stringToAccessPerm(perm string) accessPerms {
	var policy accessPerms
	switch perm {
	case "none":
		policy = accessNone
	case "readonly":
		policy = accessDownload
	case "writeonly":
		policy = accessUpload
	case "readwrite":
		policy = accessPublic
	}
	return policy
}

// doGetAccess do get access.
func doGetAccess(targetURL string) (perms accessPerms, err *probe.Error) {
	clnt, err := newClient(targetURL)
	if err != nil {
		return "", err.Trace(targetURL)
	}
	perm, err := clnt.GetAccess()
	if err != nil {
		return "", err.Trace(targetURL)
	}
	return stringToAccessPerm(perm), nil
}

// doGetAccessRules do get access rules.
func doGetAccessRules(targetURL string) (r map[string]string, err *probe.Error) {
	clnt, err := newClient(targetURL)
	if err != nil {
		return map[string]string{}, err.Trace(targetURL)
	}
	return clnt.GetAccessRules()
}

// Run policy list command
func runPolicyListCmd(ctx *cli.Context) {
	targetURL := ctx.Args().Last()
	policies, err := doGetAccessRules(targetURL)
	if err != nil {
		switch err.ToGoError().(type) {
		case APINotImplemented:
			fatalIf(err.Trace(), "Unable to list policies of a non S3 url ‘"+targetURL+"’.")
		default:
			fatalIf(err.Trace(targetURL), "Unable to list policies of target ‘"+targetURL+"’.")
		}
	}
	for k, v := range policies {
		printMsg(policyRules{Resource: k, Allow: v})
	}
}

// Run policy links command
func runPolicyLinksCmd(ctx *cli.Context) {
	// Get alias/bucket/prefix argument
	targetURL := ctx.Args().Last()

	// Fetch all policies associated to the passed url
	policies, err := doGetAccessRules(targetURL)
	if err != nil {
		switch err.ToGoError().(type) {
		case APINotImplemented:
			fatalIf(err.Trace(), "Unable to list policies of a non S3 url ‘"+targetURL+"’.")
		default:
			fatalIf(err.Trace(targetURL), "Unable to list policies of target ‘"+targetURL+"’.")
		}
	}

	// Extract alias from the passed argument, we'll need it to
	// construct new pathes to list public objects
	alias, path := url2Alias(targetURL)

	isRecursive := ctx.Bool("recursive")
	isIncomplete := false

	// Iterate over policy rules to fetch public urls, then search
	// for objects under those urls
	for k, v := range policies {
		// Trim the asterisk in policy rules
		policyPath := strings.TrimSuffix(k, "*")
		// Check if current policy prefix is related to the url passed by the user
		if !strings.HasPrefix(policyPath, path) {
			continue
		}
		// Check if the found policy has read permission
		perm := stringToAccessPerm(v)
		if perm != accessDownload && perm != accessPublic {
			continue
		}
		// Construct the new path to search for public objects
		newURL := alias + "/" + policyPath
		clnt, err := newClient(newURL)
		fatalIf(err.Trace(newURL), "Unable to initialize target ‘"+targetURL+"’.")
		// Search for public objects
		for content := range clnt.List(isRecursive, isIncomplete, DirFirst) {
			if content.Err != nil {
				errorIf(content.Err.Trace(clnt.GetURL().String()), "Unable to list folder.")
				continue
			}

			// Encode public URL
			u, e := url.Parse(content.URL.String())
			errorIf(probe.NewError(e), "Unable to parse url `"+content.URL.String()+"`")
			publicURL := u.String()

			// Construct the message to be displayed to the user
			msg := policyLinksMessage{
				Status: "success",
				URL:    publicURL,
			}
			// Print the found object
			printMsg(msg)
		}
	}
}

// Run policy cmd to fetch set permission
func runPolicyCmd(ctx *cli.Context) {
	perms := accessPerms(ctx.Args().First())
	if perms.isValidAccessPERM() {
		targetURL := ctx.Args().Last()
		err := doSetAccess(targetURL, perms)
		// Upon error exit.
		if err != nil {
			switch err.ToGoError().(type) {
			case APINotImplemented:
				fatalIf(err.Trace(), "Unable to set policy of a non S3 url ‘"+targetURL+"’.")
			default:
				fatalIf(err.Trace(targetURL, string(perms)),
					"Unable to set policy ‘"+string(perms)+"’ for ‘"+targetURL+"’.")

			}
		}

		printMsg(policyMessage{
			Status:    "success",
			Operation: "set",
			Bucket:    targetURL,
			Perms:     perms,
		})
	} else {
		targetURL := ctx.Args().First()
		perms, err := doGetAccess(targetURL)
		if err != nil {
			switch err.ToGoError().(type) {
			case APINotImplemented:
				fatalIf(err.Trace(), "Unable to get policy of a non S3 url ‘"+targetURL+"’.")
			default:
				fatalIf(err.Trace(targetURL), "Unable to get policy for ‘"+targetURL+"’.")
			}
		}

		printMsg(policyMessage{
			Status:    "success",
			Operation: "get",
			Bucket:    targetURL,
			Perms:     perms,
		})
	}
}

func mainPolicy(ctx *cli.Context) error {
	// Set global flags from context.
	setGlobalsFromContext(ctx)

	// check 'policy' cli arguments.
	checkPolicySyntax(ctx)

	// Additional command speific theme customization.
	console.SetColor("Policy", color.New(color.FgGreen, color.Bold))

	switch ctx.Args().First() {
	case "list":
		// policy list alias/bucket/prefix
		runPolicyListCmd(ctx)
	case "links":
		// policy links alias/bucket/prefix
		runPolicyLinksCmd(ctx)
	default:
		// policy [download|upload|public|] alias/bucket/prefix
		runPolicyCmd(ctx)
	}
	return nil
}
<|endoftext|>


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

package main

import (
	"flag"
	"fmt"
	"log"
	"time"

	"github.com/codeignition/recon/cmd/recond/config"
	"github.com/codeignition/recon/policy"
	"github.com/nats-io/nats"
)

const agentsAPIPath = "/api/agents" // agents path in the marksman server

// natsEncConn is the opened with the URL obtained from marksman.
// It is populated if the agent registers successfully.
var natsEncConn *nats.EncodedConn

// updateInterval is time.Duration that specifies the interval
// between two consecutive updates.
const updateInterval = 5 * time.Second

func main() {
	log.SetPrefix("recond: ")

	var marksmanAddr = flag.String("marksman", "http://localhost:3000", "address of the marksman server")
	flag.Parse()

	conf, err := config.Init()
	if err != nil {
		log.Fatalln(err)
	}

	// agent represents a single agent on which the recond
	// is running.
	var agent = &Agent{
		UID: conf.UID,
	}

	err = agent.register(*marksmanAddr)
	if err != nil {
		log.Fatalln(err)
	}

	defer natsEncConn.Close()

	natsEncConn.Subscribe(agent.UID, func(s string) {
		fmt.Printf("Received a message: %s\n", s)
	})

	natsEncConn.Subscribe(agent.UID+"_policy", func(subj, reply string, p *policy.Policy) {
		fmt.Printf("Received a Policy: %v\n", p)
		if err := conf.AddPolicy(*p); err != nil {
			natsEncConn.Publish(reply, err.Error())
			return
		}
		if err := conf.Save(); err != nil {
			natsEncConn.Publish(reply, err.Error())
			return
		}
		if err := p.Execute(); err != nil {
			natsEncConn.Publish(reply, err.Error())
			return
		}
		natsEncConn.Publish(reply, "policy ack") // acknowledge
	})

	c := time.Tick(updateInterval)
	for now := range c {
		log.Println("Update sent at", now)
		agent.update()
	}
}
<commit_msg>run stored policies on start, fix #2<commit_after>// Copyright 2015 CodeIgnition. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package main

import (
	"flag"
	"fmt"
	"log"
	"time"

	"github.com/codeignition/recon/cmd/recond/config"
	"github.com/codeignition/recon/policy"
	"github.com/nats-io/nats"
)

const agentsAPIPath = "/api/agents" // agents path in the marksman server

// natsEncConn is the opened with the URL obtained from marksman.
// It is populated if the agent registers successfully.
var natsEncConn *nats.EncodedConn

// updateInterval is time.Duration that specifies the interval
// between two consecutive updates.
const updateInterval = 5 * time.Second

func main() {
	log.SetPrefix("recond: ")

	var marksmanAddr = flag.String("marksman", "http://localhost:3000", "address of the marksman server")
	flag.Parse()

	conf, err := config.Init()
	if err != nil {
		log.Fatalln(err)
	}

	// agent represents a single agent on which the recond
	// is running.
	var agent = &Agent{
		UID: conf.UID,
	}

	err = agent.register(*marksmanAddr)
	if err != nil {
		log.Fatalln(err)
	}

	defer natsEncConn.Close()

	go runStoredPolicies(conf)

	natsEncConn.Subscribe(agent.UID, func(s string) {
		fmt.Printf("Received a message: %s\n", s)
	})

	natsEncConn.Subscribe(agent.UID+"_policy", func(subj, reply string, p *policy.Policy) {
		fmt.Printf("Received a Policy: %v\n", p)
		if err := conf.AddPolicy(*p); err != nil {
			natsEncConn.Publish(reply, err.Error())
			return
		}
		if err := conf.Save(); err != nil {
			natsEncConn.Publish(reply, err.Error())
			return
		}
		if err := p.Execute(); err != nil {
			natsEncConn.Publish(reply, err.Error())
			return
		}
		natsEncConn.Publish(reply, "policy ack") // acknowledge
	})

	c := time.Tick(updateInterval)
	for now := range c {
		log.Println("Update sent at", now)
		agent.update()
	}
}

func runStoredPolicies(c *config.Config) {
	for _, p := range c.PolicyConfig {
		go func() {
			if err := p.Execute(); err != nil {
				log.Fatal(err)
			}
		}()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"strings"

	"github.com/eahydra/socks"
)

func main() {
	confGroup, err := LoadConfigGroup("socks.config")
	if err != nil {
		ErrLog.Println("initGlobalConfig failed, err:", err)
		return
	}
	InfoLog.Println(confGroup)

	for _, conf := range confGroup.AllConfig {
		router := BuildUpstreamRouter(conf)
		runHTTPProxyServer(conf, router)
		runSOCKS4Server(conf, router)
		runSOCKS5Server(conf, router)
	}
	go http.ListenAndServe(confGroup.PprofAddr, nil)

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Kill, os.Interrupt)
	<-sigChan
}

func BuildUpstreamRouter(conf Config) socks.Dialer {
	var allForward []socks.Dialer
	for _, upstreamConf := range conf.AllUpstreamConfig {
		var forward socks.Dialer
		forward = NewDecorateDirect(conf.DNSCacheTimeout)
		cipherDecorator := NewCipherConnDecorator(upstreamConf.CryptoMethod, upstreamConf.Password)
		forward = NewDecorateClient(forward, cipherDecorator)

		var err error
		switch strings.ToLower(upstreamConf.ServerType) {
		case "socks5":
			{
				forward, err = socks.NewSocks5Client("tcp", upstreamConf.Addr, "", "", forward)
			}
		case "shadowsocks":
			{
				forward, err = socks.NewShadowSocksClient("tcp", upstreamConf.Addr, forward)
			}
		}
		if err != nil {
			ErrLog.Println("build upstream failed, err:", err, upstreamConf.ServerType, upstreamConf.Addr)
			continue
		}
		allForward = append(allForward, forward)
	}
	if len(allForward) == 0 {
		router := NewDecorateDirect(conf.DNSCacheTimeout)
		allForward = append(allForward, router)
	}
	return NewUpstreamDialer(allForward)
}

func runHTTPProxyServer(conf Config, router socks.Dialer) {
	if conf.HTTPProxyAddr != "" {
		listener, err := net.Listen("tcp", conf.HTTPProxyAddr)
		if err != nil {
			ErrLog.Println("net.Listen at ", conf.HTTPProxyAddr, " failed, err:", err)
			return
		}
		go func() {
			defer listener.Close()
			httpProxy := socks.NewHTTPProxy(router)
			http.Serve(listener, httpProxy)
		}()
	}
}

func runSOCKS4Server(conf Config, forward socks.Dialer) {
	if conf.SOCKS4Addr != "" {
		listener, err := net.Listen("tcp", conf.SOCKS4Addr)
		if err != nil {
			ErrLog.Println("net.Listen failed, err:", err, conf.SOCKS4Addr)
			return
		}
		cipherDecorator := NewCipherConnDecorator(conf.LocalCryptoMethod, conf.LocalCryptoPassword)
		listener = NewDecorateListener(listener, cipherDecorator)
		socks4Svr, err := socks.NewSocks4Server(forward)
		if err != nil {
			listener.Close()
			ErrLog.Println("socks.NewSocks4Server failed, err:", err)
		}
		go func() {
			defer listener.Close()
			socks4Svr.Serve(listener)
		}()
	}
}

func runSOCKS5Server(conf Config, forward socks.Dialer) {
	if conf.SOCKS5Addr != "" {
		listener, err := net.Listen("tcp", conf.SOCKS5Addr)
		if err != nil {
			ErrLog.Println("net.Listen failed, err:", err, conf.SOCKS5Addr)
			return
		}
		cipherDecorator := NewCipherConnDecorator(conf.LocalCryptoMethod, conf.LocalCryptoPassword)
		listener = NewDecorateListener(listener, cipherDecorator)
		socks5Svr, err := socks.NewSocks5Server(forward)
		if err != nil {
			listener.Close()
			ErrLog.Println("socks.NewSocks5Server failed, err:", err)
			return
		}
		go func() {
			defer listener.Close()
			socks5Svr.Serve(listener)
		}()
	}
}
<commit_msg>support -c command param to special config file.<commit_after>package main

import (
	"flag"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"strings"

	"github.com/eahydra/socks"
)

func main() {
	var configFile string
	flag.StringVar(&configFile, "c", "socks.config", "config file path")
	flag.Parse()

	confGroup, err := LoadConfigGroup(configFile)
	if err != nil {
		ErrLog.Println("initGlobalConfig failed, err:", err)
		return
	}
	InfoLog.Println(confGroup)

	for _, conf := range confGroup.AllConfig {
		router := BuildUpstreamRouter(conf)
		runHTTPProxyServer(conf, router)
		runSOCKS4Server(conf, router)
		runSOCKS5Server(conf, router)
	}
	go http.ListenAndServe(confGroup.PprofAddr, nil)

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Kill, os.Interrupt)
	<-sigChan
}

func BuildUpstreamRouter(conf Config) socks.Dialer {
	var allForward []socks.Dialer
	for _, upstreamConf := range conf.AllUpstreamConfig {
		var forward socks.Dialer
		forward = NewDecorateDirect(conf.DNSCacheTimeout)
		cipherDecorator := NewCipherConnDecorator(upstreamConf.CryptoMethod, upstreamConf.Password)
		forward = NewDecorateClient(forward, cipherDecorator)

		var err error
		switch strings.ToLower(upstreamConf.ServerType) {
		case "socks5":
			{
				forward, err = socks.NewSocks5Client("tcp", upstreamConf.Addr, "", "", forward)
			}
		case "shadowsocks":
			{
				forward, err = socks.NewShadowSocksClient("tcp", upstreamConf.Addr, forward)
			}
		}
		if err != nil {
			ErrLog.Println("build upstream failed, err:", err, upstreamConf.ServerType, upstreamConf.Addr)
			continue
		}
		allForward = append(allForward, forward)
	}
	if len(allForward) == 0 {
		router := NewDecorateDirect(conf.DNSCacheTimeout)
		allForward = append(allForward, router)
	}
	return NewUpstreamDialer(allForward)
}

func runHTTPProxyServer(conf Config, router socks.Dialer) {
	if conf.HTTPProxyAddr != "" {
		listener, err := net.Listen("tcp", conf.HTTPProxyAddr)
		if err != nil {
			ErrLog.Println("net.Listen at ", conf.HTTPProxyAddr, " failed, err:", err)
			return
		}
		go func() {
			defer listener.Close()
			httpProxy := socks.NewHTTPProxy(router)
			http.Serve(listener, httpProxy)
		}()
	}
}

func runSOCKS4Server(conf Config, forward socks.Dialer) {
	if conf.SOCKS4Addr != "" {
		listener, err := net.Listen("tcp", conf.SOCKS4Addr)
		if err != nil {
			ErrLog.Println("net.Listen failed, err:", err, conf.SOCKS4Addr)
			return
		}
		cipherDecorator := NewCipherConnDecorator(conf.LocalCryptoMethod, conf.LocalCryptoPassword)
		listener = NewDecorateListener(listener, cipherDecorator)
		socks4Svr, err := socks.NewSocks4Server(forward)
		if err != nil {
			listener.Close()
			ErrLog.Println("socks.NewSocks4Server failed, err:", err)
		}
		go func() {
			defer listener.Close()
			socks4Svr.Serve(listener)
		}()
	}
}

func runSOCKS5Server(conf Config, forward socks.Dialer) {
	if conf.SOCKS5Addr != "" {
		listener, err := net.Listen("tcp", conf.SOCKS5Addr)
		if err != nil {
			ErrLog.Println("net.Listen failed, err:", err, conf.SOCKS5Addr)
			return
		}
		cipherDecorator := NewCipherConnDecorator(conf.LocalCryptoMethod, conf.LocalCryptoPassword)
		listener = NewDecorateListener(listener, cipherDecorator)
		socks5Svr, err := socks.NewSocks5Server(forward)
		if err != nil {
			listener.Close()
			ErrLog.Println("socks.NewSocks5Server failed, err:", err)
			return
		}
		go func() {
			defer listener.Close()
			socks5Svr.Serve(listener)
		}()
	}
}
<|endoftext|>


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

package filesystem

import (
	"fmt"
	"github.com/andreaskoch/allmark2/common/config"
	"github.com/andreaskoch/allmark2/common/logger"
	"github.com/andreaskoch/allmark2/common/route"
	"github.com/andreaskoch/allmark2/common/util/fsutil"
	"github.com/andreaskoch/allmark2/dataaccess"
	"path/filepath"
	"strings"
)

type Repository struct {
	logger    logger.Logger
	hash      string
	directory string
}

func NewRepository(logger logger.Logger, directory string) (*Repository, error) {

	// check if path exists
	if !fsutil.PathExists(directory) {
		return nil, fmt.Errorf("The path %q does not exist.", directory)
	}

	// check if the supplied path is a file
	if isDirectory, _ := fsutil.IsDirectory(directory); !isDirectory {
		directory = filepath.Dir(directory)
	}

	// abort if the supplied path is a reserved directory
	if isReservedDirectory(directory) {
		return nil, fmt.Errorf("The path %q is using a reserved name and cannot be a root.", directory)
	}

	// hash provider: use the directory name for the hash (for now)
	directoryName := strings.ToLower(filepath.Base(directory))
	hash, err := getStringHash(directoryName)
	if err != nil {
		return nil, fmt.Errorf("Cannot create a hash for the repository with the name %q. Error: %s", directoryName, err)
	}

	return &Repository{
		logger:    logger,
		directory: directory,
		hash:      hash,
	}, nil
}

func (repository *Repository) GetItems() (itemEvents chan *dataaccess.RepositoryEvent, done chan bool) {

	itemEvents = make(chan *dataaccess.RepositoryEvent, 1)
	done = make(chan bool)

	go func() {

		// repository directory item
		indexItems(repository, repository.directory, itemEvents)

		done <- true
	}()

	return itemEvents, done
}

func (repository *Repository) Id() string {
	return repository.hash
}

func (repository *Repository) Path() string {
	return repository.directory
}

// Create a new Item for the specified path.
func indexItems(repository *Repository, itemPath string, itemEvents chan *dataaccess.RepositoryEvent) {

	// abort if path does not exist
	if !fsutil.PathExists(itemPath) {
		itemEvents <- dataaccess.NewEvent(nil, fmt.Errorf("The path %q does not exist.", itemPath))
		return
	}

	// abort if path is reserved
	if isReservedDirectory(itemPath) {
		itemEvents <- dataaccess.NewEvent(nil, fmt.Errorf("The path %q is using a reserved name and cannot be an item.", itemPath))
		return
	}

	// check if its a virtual item or a markdown item
	itemDirectory := filepath.Dir(itemPath)
	if isDirectory, _ := fsutil.IsDirectory(itemPath); isDirectory {

		if found, filepath := findMarkdownFileInDirectory(itemPath); found {

			itemDirectory = itemPath
			itemPath = filepath

		} else {

			// virtual item
			itemDirectory = itemPath

		}

	} else if !isMarkdownFile(itemPath) {

		// the supplied item path does not point to a markdown file
		itemEvents <- dataaccess.NewEvent(nil, fmt.Errorf("%q is not a markdown file.", itemPath))
		return
	}

	// route
	route, err := route.New(repository.Path(), itemPath)
	if err != nil {
		itemEvents <- dataaccess.NewEvent(nil, fmt.Errorf("Cannot create an Item for the path %q. Error: %s", itemPath, err))
	}

	// content provider
	contentProvider := newContentProvider(itemPath, route)

	// create the file index
	filesDirectory := filepath.Join(itemDirectory, config.FilesDirectoryName)
	files := getFiles(repository, filesDirectory)

	// create the item
	item, err := dataaccess.NewItem(route, contentProvider, files)

	itemEvents <- dataaccess.NewEvent(item, err)

	// recurse for child items
	childItemDirectories := getChildDirectories(itemDirectory)
	for _, childItemDirectory := range childItemDirectories {
		indexItems(repository, childItemDirectory, itemEvents)
	}
}
<commit_msg>Refactored the repository indexing. Don't index directories as "virtual" items. A route to a file and an Item should always have a "file" behind it. Later I might implement virtual items with a "file-emulation" ...<commit_after>// Copyright 2013 Andreas Koch. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package filesystem

import (
	"fmt"
	"github.com/andreaskoch/allmark2/common/config"
	"github.com/andreaskoch/allmark2/common/logger"
	"github.com/andreaskoch/allmark2/common/route"
	"github.com/andreaskoch/allmark2/common/util/fsutil"
	"github.com/andreaskoch/allmark2/dataaccess"
	"path/filepath"
	"strings"
)

type Repository struct {
	logger    logger.Logger
	hash      string
	directory string
}

func NewRepository(logger logger.Logger, directory string) (*Repository, error) {

	// check if path exists
	if !fsutil.PathExists(directory) {
		return nil, fmt.Errorf("The path %q does not exist.", directory)
	}

	// check if the supplied path is a file
	if isDirectory, _ := fsutil.IsDirectory(directory); !isDirectory {
		directory = filepath.Dir(directory)
	}

	// abort if the supplied path is a reserved directory
	if isReservedDirectory(directory) {
		return nil, fmt.Errorf("The path %q is using a reserved name and cannot be a root.", directory)
	}

	// hash provider: use the directory name for the hash (for now)
	directoryName := strings.ToLower(filepath.Base(directory))
	hash, err := getStringHash(directoryName)
	if err != nil {
		return nil, fmt.Errorf("Cannot create a hash for the repository with the name %q. Error: %s", directoryName, err)
	}

	return &Repository{
		logger:    logger,
		directory: directory,
		hash:      hash,
	}, nil
}

func (repository *Repository) GetItems() (itemEvents chan *dataaccess.RepositoryEvent, done chan bool) {

	itemEvents = make(chan *dataaccess.RepositoryEvent, 1)
	done = make(chan bool)

	go func() {

		// repository directory item
		indexItems(repository, repository.directory, itemEvents)

		done <- true
	}()

	return itemEvents, done
}

func (repository *Repository) Id() string {
	return repository.hash
}

func (repository *Repository) Path() string {
	return repository.directory
}

// Create a new Item for the specified path.
func indexItems(repository *Repository, itemPath string, itemEvents chan *dataaccess.RepositoryEvent) {

	// abort if path does not exist
	if !fsutil.PathExists(itemPath) {
		itemEvents <- dataaccess.NewEvent(nil, fmt.Errorf("The path %q does not exist.", itemPath))
		return
	}

	// abort if path is reserved
	if isReservedDirectory(itemPath) {
		itemEvents <- dataaccess.NewEvent(nil, fmt.Errorf("The path %q is using a reserved name and cannot be an item.", itemPath))
		return
	}

	// make sure the item path points to a markdown file
	isVirtualItem := false
	itemDirectory := filepath.Dir(itemPath)
	if isDirectory, _ := fsutil.IsDirectory(itemPath); isDirectory {

		// search for a markdown file in the directory
		if found, filepath := findMarkdownFileInDirectory(itemPath); found {

			itemDirectory = itemPath
			itemPath = filepath

		} else {

			// virtual item
			isVirtualItem = true
			itemDirectory = itemPath

		}

	} else if !isMarkdownFile(itemPath) {

		// the supplied item path does not point to a markdown file
		itemEvents <- dataaccess.NewEvent(nil, fmt.Errorf("%q is not a markdown file.", itemPath))
		return
	}

	// create a new item
	if !isVirtualItem {
		// route
		route, err := route.New(repository.Path(), itemPath)
		if err != nil {
			itemEvents <- dataaccess.NewEvent(nil, fmt.Errorf("Cannot create an Item for the path %q. Error: %s", itemPath, err))
		}

		// content provider
		contentProvider := newContentProvider(itemPath, route)

		// create the file index
		filesDirectory := filepath.Join(itemDirectory, config.FilesDirectoryName)
		files := getFiles(repository, filesDirectory)

		// create the item
		item, err := dataaccess.NewItem(route, contentProvider, files)

		itemEvents <- dataaccess.NewEvent(item, err)
	}

	// recurse for child items
	childItemDirectories := getChildDirectories(itemDirectory)
	for _, childItemDirectory := range childItemDirectories {
		indexItems(repository, childItemDirectory, itemEvents)
	}
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"fmt"
	"github.com/szabba/md/newton"
	"github.com/szabba/md/vect"
	"io"
)

// A force that is always zero
type ZeroForce struct{}

func (_ ZeroForce) Accel(bs []*newton.Body, i int) (a vect.Vector) {

	return vect.Zero
}

// A 'picky' force, that doesn't affect some bodies
type PickyForce struct {
	force   newton.Force
	zeroFor []int
}

// Creates a picky version of a force
func NewPicky(f newton.Force, zeroFor ...int) newton.Force {

	return &PickyForce{force: f, zeroFor: zeroFor}
}

func (picky *PickyForce) Accel(bs []*newton.Body, i int) (a vect.Vector) {

	for _, ignored := range picky.zeroFor {

		if ignored == i {

			return vect.Zero
		}
	}

	return picky.force.Accel(bs, i)
}

type ParticleRect struct {
	*newton.System
	rows, cols int
}

// Creates a rectangular grid of particles
func NewRect(rows, cols int) *ParticleRect {

	rect := &ParticleRect{
		rows: rows, cols: cols,
	}

	rect.System = newton.NewSystem(newton.Verlet, rows*cols)

	for i := 0; i < rect.Bodies(); i++ {

		b := rect.Body(i)

		b.SetMass(1)

		pos := rect.RestingPosition(i)

		b.Shift(pos, vect.Zero)
		b.Shift(pos, vect.Zero)

	}

	rect.SetForce(ZeroForce{})

	return rect
}

// The row and column in which the i-th particle is
func (rect *ParticleRect) RowAndColumn(ith int) (row, col int) {

	return ith % rect.rows, ith / rect.rows
}

// Initial, resting position of the i-th particle
func (rect *ParticleRect) RestingPosition(ith int) vect.Vector {

	row, col := rect.RowAndColumn(ith)

	return vect.UnitX.Scale(float64(row)).Plus(
		vect.UnitX.Scale(float64(col)),
	)
}

// Dimmensions of the rectangle
func (rect *ParticleRect) Size() (rows, cols int) {

	return rect.rows, rect.cols
}

// An output formatting type
type Formatter struct {
	rect *ParticleRect
}

// Formats a data header
func (f Formatter) Header(writeTo io.Writer) {

	fmt.Fprintf(writeTo, "%d\n\n", f.rect.Bodies())
}

// Formats the description of ball states
func (f Formatter) Frame(writeTo io.Writer) {

	for i := 0; i < f.rect.Bodies(); i++ {

		b := f.rect.Body(i)

		x, v := b.Now()

		fmt.Fprintf(
			writeTo, "%d %f %f %f %f %f %f\n", i,
			x[0], x[1], x[2],
			v[0], v[1], v[2],
		)

	}
	fmt.Fprintf(writeTo, "\n")
}

func main() {

	rect := NewRect(20, 20)
	rect.Step(0.05)
}
<commit_msg>Let a Formatter remember where to write to<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"fmt"
	"github.com/szabba/md/newton"
	"github.com/szabba/md/vect"
	"io"
)

// A force that is always zero
type ZeroForce struct{}

func (_ ZeroForce) Accel(bs []*newton.Body, i int) (a vect.Vector) {

	return vect.Zero
}

// A 'picky' force, that doesn't affect some bodies
type PickyForce struct {
	force   newton.Force
	zeroFor []int
}

// Creates a picky version of a force
func NewPicky(f newton.Force, zeroFor ...int) newton.Force {

	return &PickyForce{force: f, zeroFor: zeroFor}
}

func (picky *PickyForce) Accel(bs []*newton.Body, i int) (a vect.Vector) {

	for _, ignored := range picky.zeroFor {

		if ignored == i {

			return vect.Zero
		}
	}

	return picky.force.Accel(bs, i)
}

type ParticleRect struct {
	*newton.System
	rows, cols int
}

// Creates a rectangular grid of particles
func NewRect(rows, cols int) *ParticleRect {

	rect := &ParticleRect{
		rows: rows, cols: cols,
	}

	rect.System = newton.NewSystem(newton.Verlet, rows*cols)

	for i := 0; i < rect.Bodies(); i++ {

		b := rect.Body(i)

		b.SetMass(1)

		pos := rect.RestingPosition(i)

		b.Shift(pos, vect.Zero)
		b.Shift(pos, vect.Zero)

	}

	rect.SetForce(ZeroForce{})

	return rect
}

// The row and column in which the i-th particle is
func (rect *ParticleRect) RowAndColumn(ith int) (row, col int) {

	return ith % rect.rows, ith / rect.rows
}

// Initial, resting position of the i-th particle
func (rect *ParticleRect) RestingPosition(ith int) vect.Vector {

	row, col := rect.RowAndColumn(ith)

	return vect.UnitX.Scale(float64(row)).Plus(
		vect.UnitX.Scale(float64(col)),
	)
}

// Dimmensions of the rectangle
func (rect *ParticleRect) Size() (rows, cols int) {

	return rect.rows, rect.cols
}

// An output formatting type
type Formatter struct {
	rect    *ParticleRect
	writeTo io.Writer
}

// Formats a data header
func (f Formatter) Header() {

	fmt.Fprintf(f.writeTo, "%d\n\n", f.rect.Bodies())
}

// Formats the description of ball states
func (f Formatter) Frame() {

	for i := 0; i < f.rect.Bodies(); i++ {

		b := f.rect.Body(i)

		x, v := b.Now()

		fmt.Fprintf(
			f.writeTo, "%d %f %f %f %f %f %f\n", i,
			x[0], x[1], x[2],
			v[0], v[1], v[2],
		)

	}
	fmt.Fprintf(f.writeTo, "\n")
}

func main() {

	rect := NewRect(20, 20)
	rect.Step(0.05)
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"flag"
	"fmt"
	"github.com/szabba/md/newton"
	"github.com/szabba/md/vect"
	"io"
	"log"
	"os"
	"strconv"
)

// A constant force
type ConstForce vect.Vector

func (f ConstForce) Accel(bs []*newton.Body, i int) (a vect.Vector) {

	return vect.Vector(f).Scale(1 / bs[i].Mass())
}

// A 'picky' force, that doesn't affect some bodies
type PickyForce struct {
	force   newton.Force
	zeroFor []int
}

// Creates a picky version of a force
func NewPicky(f newton.Force, zeroFor ...int) newton.Force {

	return &PickyForce{force: f, zeroFor: zeroFor}
}

func (picky *PickyForce) Accel(bs []*newton.Body, i int) (a vect.Vector) {

	for _, ignored := range picky.zeroFor {

		if ignored == i {

			return vect.Zero
		}
	}

	return picky.force.Accel(bs, i)
}

type ParticleRect struct {
	*newton.System
	rows, cols int
}

// Creates a rectangular grid of particles
//
// The program's behaviour when rows or cols are less than 1 is unspecified
func NewRect(rows, cols int) *ParticleRect {

	rect := &ParticleRect{
		rows: rows, cols: cols,
	}

	rect.System = newton.NewSystem(newton.Verlet, rows*cols)

	for i := 0; i < rect.Bodies(); i++ {

		b := rect.Body(i)

		b.SetMass(1)

		pos := rect.RestingPosition(i)

		b.Shift(pos, vect.Zero)
		b.Shift(pos, vect.Zero)

	}

	rect.SetForce(ConstForce(vect.Zero))

	return rect
}

// The row and column in which the i-th particle is
func (rect *ParticleRect) RowAndColumn(ith int) (row, col int) {

	return ith % rect.rows, ith / rect.rows
}

// Initial, resting position of the i-th particle
func (rect *ParticleRect) RestingPosition(ith int) vect.Vector {

	row, col := rect.RowAndColumn(ith)

	return vect.UnitX.Scale(float64(row)).Plus(
		vect.UnitY.Scale(float64(col)),
	)
}

// Dimmensions of the rectangle
func (rect *ParticleRect) Size() (rows, cols int) {

	return rect.rows, rect.cols
}

// Are the i-th and j-th particles neighbours?
func (rect *ParticleRect) Neighbours(i, j int) bool {

	xI, yI := rect.RowAndColumn(i)
	xJ, yJ := rect.RowAndColumn(j)

	nearInX := xI-1 == xJ || xJ == xI+1
	nearInY := yI-1 == yJ || yJ == yI+1

	sameX := xI == xJ
	sameY := yI == yJ

	colNeighbour := sameX && nearInY
	rowNeighbour := sameY && nearInX

	return colNeighbour || rowNeighbour
}

// Prepare a Hooke's force bidning neihbouring particles
func (rect *ParticleRect) Hooke(k float64) newton.Force {

	var h newton.Hooke

	h.Springs = make([][]newton.Spring, rect.Bodies())
	for i, _ := range h.Springs {

		h.Springs[i] = make([]newton.Spring, rect.Bodies())
		for j, _ := range h.Springs[i] {

			if rect.Neighbours(i, j) {

				h.Springs[i][j].K = k

			}
			h.Springs[i][j].L0 = 1
		}
	}

	return h
}

// Prepare a constant force that only affect the center of the rectangle
//
// Depending on the shape of the rectangle this will pull 1, 2 or 4 particles.
// The force applied per particle will be divided by this number.
func (rect *ParticleRect) CentralPull(pull vect.Vector) newton.Force {

	rows, cols := rect.Size()

	xRange, yRange := make([]int, rows%2), make([]int, cols%2)
	for i, x := 0, rows/2; i < len(xRange); i, x = i+1, x+1 {

		xRange[i] = x
	}
	for i, y := 0, cols/2; i < len(yRange); i, y = i+1, y+1 {

		yRange[i] = y
	}

	i, picked := 0, make([]int, len(xRange)*len(yRange))
	for _, x := range xRange {
		for _, y := range yRange {

			picked[i] = y*cols + x
			i++
		}
	}

	return NewPicky(
		ConstForce(pull.Scale(1/float64(len(picked)))),
		picked...,
	)
}

// Runs the simulation for the given number of steps at a time step of dt
// printing to writeTo
func (rect *ParticleRect) Run(writeTo io.Writer, dt float64, steps int) {

	format := &Formatter{rect: rect, writeTo: writeTo}

	format.Header()

	for i := 0; i < steps; i++ {

		format.Frame()

		rect.Step(dt)
	}
}

// An output formatting type
type Formatter struct {
	rect    *ParticleRect
	writeTo io.Writer
}

// Formats a data header
func (f Formatter) Header() {

	fmt.Fprintf(f.writeTo, "%d\n\n", f.rect.Bodies())
}

// Formats the description of ball states
func (f Formatter) Frame() {

	for i := 0; i < f.rect.Bodies(); i++ {

		b := f.rect.Body(i)

		x, v := b.Now()

		fmt.Fprintf(
			f.writeTo, "%d %f %f %f %f %f %f\n", i,
			x[0], x[1], x[2],
			v[0], v[1], v[2],
		)

	}
	fmt.Fprintf(f.writeTo, "\n")
}

const usage string = `Usage of %s:

	Simulate a rectangular surface made of particles interconnected with
	strings.

		%s [options] ROWS COLS

	ROWS and COLS determines the shape and size of the rectangle. They must
	both be at least 1.

Options:
`

func Help() {

	program := os.Args[0]

	fmt.Fprintf(os.Stderr, usage, program, program)

	flag.PrintDefaults()
}

func main() {

	var (
		usage    bool
		p, k, dt float64
		steps    int
	)

	log.SetFlags(0)

	flag.Float64Var(
		&p, "pull", 1,
		"Magnitude of the vertical pulling force. When negative, the force pulls down.",
	)
	flag.Float64Var(&dt, "dt", 0.05, "Time step")
	flag.Float64Var(&k, "k", 1, "Hooke's constant")
	flag.IntVar(&steps, "steps", 5, "Simulation steps to perform")
	flag.BoolVar(&usage, "help", false, "Print usage string")

	flag.Parse()

	if usage {

		Help()

	} else if len(flag.Args()) != 2 {

		log.Fatal("Both ROWS and COLS need to be specified")

	} else {

		var (
			rows, cols int
			err        error
		)

		rows, err = strconv.Atoi(flag.Arg(0))
		if err != nil {

			log.Print("ROWS needs to be an integer")
			log.Fatal(err.Error())
		}
		cols, err = strconv.Atoi(flag.Arg(1))
		if err != nil {

			log.Print("COLS needs to be an integer")
			log.Fatal(err.Error())
		}

		rect := NewRect(rows, cols)
		rect.AddForce(rect.Hooke(k))
		rect.AddForce(rect.CentralPull(vect.UnitZ.Scale(p)))
		rect.Run(os.Stdout, dt, steps)
	}
}
<commit_msg>Fix ParticleRect.CentralPull<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"flag"
	"fmt"
	"github.com/szabba/md/newton"
	"github.com/szabba/md/vect"
	"io"
	"log"
	"os"
	"strconv"
)

// A constant force
type ConstForce vect.Vector

func (f ConstForce) Accel(bs []*newton.Body, i int) (a vect.Vector) {

	return vect.Vector(f).Scale(1 / bs[i].Mass())
}

// A 'picky' force, that doesn't affect some bodies
type PickyForce struct {
	force   newton.Force
	zeroFor []int
}

// Creates a picky version of a force
func NewPicky(f newton.Force, zeroFor ...int) newton.Force {

	return &PickyForce{force: f, zeroFor: zeroFor}
}

func (picky *PickyForce) Accel(bs []*newton.Body, i int) (a vect.Vector) {

	for _, ignored := range picky.zeroFor {

		if ignored == i {

			return vect.Zero
		}
	}

	return picky.force.Accel(bs, i)
}

type ParticleRect struct {
	*newton.System
	rows, cols int
}

// Creates a rectangular grid of particles
//
// The program's behaviour when rows or cols are less than 1 is unspecified
func NewRect(rows, cols int) *ParticleRect {

	rect := &ParticleRect{
		rows: rows, cols: cols,
	}

	rect.System = newton.NewSystem(newton.Verlet, rows*cols)

	for i := 0; i < rect.Bodies(); i++ {

		b := rect.Body(i)

		b.SetMass(1)

		pos := rect.RestingPosition(i)

		b.Shift(pos, vect.Zero)
		b.Shift(pos, vect.Zero)

	}

	rect.SetForce(ConstForce(vect.Zero))

	return rect
}

// The row and column in which the i-th particle is
func (rect *ParticleRect) RowAndColumn(ith int) (row, col int) {

	return ith % rect.rows, ith / rect.rows
}

// Initial, resting position of the i-th particle
func (rect *ParticleRect) RestingPosition(ith int) vect.Vector {

	row, col := rect.RowAndColumn(ith)

	return vect.UnitX.Scale(float64(row)).Plus(
		vect.UnitY.Scale(float64(col)),
	)
}

// Dimmensions of the rectangle
func (rect *ParticleRect) Size() (rows, cols int) {

	return rect.rows, rect.cols
}

// Are the i-th and j-th particles neighbours?
func (rect *ParticleRect) Neighbours(i, j int) bool {

	xI, yI := rect.RowAndColumn(i)
	xJ, yJ := rect.RowAndColumn(j)

	nearInX := xI-1 == xJ || xJ == xI+1
	nearInY := yI-1 == yJ || yJ == yI+1

	sameX := xI == xJ
	sameY := yI == yJ

	colNeighbour := sameX && nearInY
	rowNeighbour := sameY && nearInX

	return colNeighbour || rowNeighbour
}

// Prepare a Hooke's force bidning neihbouring particles
func (rect *ParticleRect) Hooke(k float64) newton.Force {

	var h newton.Hooke

	h.Springs = make([][]newton.Spring, rect.Bodies())
	for i, _ := range h.Springs {

		h.Springs[i] = make([]newton.Spring, rect.Bodies())
		for j, _ := range h.Springs[i] {

			if rect.Neighbours(i, j) {

				h.Springs[i][j].K = k

			}
			h.Springs[i][j].L0 = 1
		}
	}

	return h
}

// Is the i-th particle near the center in it's resting position?
func (rect *ParticleRect) NearCenter(ith int) bool {

	rows, cols := rect.Size()

	center := vect.NewVector(float64(rows)/2, float64(cols)/2, 0)

	fromCenter := rect.RestingPosition(ith).Minus(center)

	return fromCenter.Norm() < 1
}

// Number of particles that are not near the center in their resting positions
func (rect *ParticleRect) ExceptCenter() []int {

	rows, cols := rect.Size()

	centerSize := 1
	if rows%2 == 0 {
		centerSize *= 2
	}
	if cols%2 == 0 {
		centerSize *= 2
	}

	rest := make([]int, rows*cols-centerSize)

	for i, j := 0, 0; i < len(rest); j++ {

		if !rect.NearCenter(j) {

			i, rest[i] = i+1, j
		}
	}

	return rest
}

// Prepare a constant force that only affects the center of the rectangle (ie
// ignores everything except it)
//
// Depending on the shape of the rectangle this will pull 1, 2 or 4 particles.
// The force applied per particle will be divided by this number.
func (rect *ParticleRect) CentralPull(pull vect.Vector) newton.Force {

	ignored := rect.ExceptCenter()

	pulled := float64(rect.Bodies() - len(ignored))

	return NewPicky(ConstForce(pull.Scale(1/pulled)), ignored...)
}

// Runs the simulation for the given number of steps at a time step of dt
// printing to writeTo
func (rect *ParticleRect) Run(writeTo io.Writer, dt float64, steps int) {

	format := &Formatter{rect: rect, writeTo: writeTo}

	format.Header()

	for i := 0; i < steps; i++ {

		format.Frame()

		rect.Step(dt)
	}
}

// An output formatting type
type Formatter struct {
	rect    *ParticleRect
	writeTo io.Writer
}

// Formats a data header
func (f Formatter) Header() {

	fmt.Fprintf(f.writeTo, "%d\n\n", f.rect.Bodies())
}

// Formats the description of ball states
func (f Formatter) Frame() {

	for i := 0; i < f.rect.Bodies(); i++ {

		b := f.rect.Body(i)

		x, v := b.Now()

		fmt.Fprintf(
			f.writeTo, "%d %f %f %f %f %f %f\n", i,
			x[0], x[1], x[2],
			v[0], v[1], v[2],
		)

	}
	fmt.Fprintf(f.writeTo, "\n")
}

const usage string = `Usage of %s:

	Simulate a rectangular surface made of particles interconnected with
	strings.

		%s [options] ROWS COLS

	ROWS and COLS determines the shape and size of the rectangle. They must
	both be at least 1.

Options:
`

func Help() {

	program := os.Args[0]

	fmt.Fprintf(os.Stderr, usage, program, program)

	flag.PrintDefaults()
}

func main() {

	var (
		usage    bool
		p, k, dt float64
		steps    int
	)

	log.SetFlags(0)

	flag.Float64Var(
		&p, "pull", 1,
		"Magnitude of the vertical pulling force. When negative, the force pulls down.",
	)
	flag.Float64Var(&dt, "dt", 0.05, "Time step")
	flag.Float64Var(&k, "k", 1, "Hooke's constant")
	flag.IntVar(&steps, "steps", 5, "Simulation steps to perform")
	flag.BoolVar(&usage, "help", false, "Print usage string")

	flag.Parse()

	if usage {

		Help()

	} else if len(flag.Args()) != 2 {

		log.Fatal("Both ROWS and COLS need to be specified")

	} else {

		var (
			rows, cols int
			err        error
		)

		rows, err = strconv.Atoi(flag.Arg(0))
		if err != nil {

			log.Print("ROWS needs to be an integer")
			log.Fatal(err.Error())
		}
		cols, err = strconv.Atoi(flag.Arg(1))
		if err != nil {

			log.Print("COLS needs to be an integer")
			log.Fatal(err.Error())
		}

		rect := NewRect(rows, cols)
		rect.AddForce(rect.Hooke(k))
		rect.AddForce(rect.CentralPull(vect.UnitZ.Scale(p)))
		rect.Run(os.Stdout, dt, steps)
	}
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"github.com/szabba/md/newton"
	"github.com/szabba/md/vect"
)

// A 'picky' force, that doesn't affect some bodies
type PickyForce struct {
	force   newton.Force
	zeroFor []int
}

// Creates a picky version of a force
func NewPicky(f newton.Force, zeroFor ...int) newton.Force {

	return &PickyForce{force: f, zeroFor: zeroFor}
}

func (picky *PickyForce) Accel(bs []*newton.Body, i int) (a vect.Vector) {

	for _, ignored := range picky.zeroFor {

		if ignored == i {

			return vect.Zero
		}
	}

	return picky.force.Accel(bs, i)
}

type ParticleRect struct {
	*newton.System
	rows, cols int
}

// Creates a rectangular grid of particles
func NewRect(rows, cols int) *ParticleRect {

	rect := &ParticleRect{
		rows: rows, cols: cols,
	}

	rect.System = newton.NewSystem(newton.Verlet, rows*cols)

	for i := 0; i < rect.Bodies(); i++ {

		b := rect.Body(i)

		b.SetMass(1)

		pos := rect.RestingPosition(i)

		b.Shift(pos, vect.Zero)
		b.Shift(pos, vect.Zero)

	}

	return rect
}

// The row and column in which the i-th particle is
func (rect *ParticleRect) RowAndColumn(ith int) (row, col int) {

	return i % rect.rows, i / rect.rows
}

// Initial, resting position of the i-th particle
func (rect *ParticleRect) RestingPosition(ith int) vect.Vector {

	row, col := rect.RowAndColumn(ith)

	return vect.UnitX.Scale(float64(row)).Plus(
		vect.UnitX.Scale(float64(col)),
	)
}

func (rect *ParticleRect) Size() (rows, cols int) {

	return rect.rows, rect.cols
}

func main() {
}
<commit_msg>Document ParticleRect.Size<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"github.com/szabba/md/newton"
	"github.com/szabba/md/vect"
)

// A 'picky' force, that doesn't affect some bodies
type PickyForce struct {
	force   newton.Force
	zeroFor []int
}

// Creates a picky version of a force
func NewPicky(f newton.Force, zeroFor ...int) newton.Force {

	return &PickyForce{force: f, zeroFor: zeroFor}
}

func (picky *PickyForce) Accel(bs []*newton.Body, i int) (a vect.Vector) {

	for _, ignored := range picky.zeroFor {

		if ignored == i {

			return vect.Zero
		}
	}

	return picky.force.Accel(bs, i)
}

type ParticleRect struct {
	*newton.System
	rows, cols int
}

// Creates a rectangular grid of particles
func NewRect(rows, cols int) *ParticleRect {

	rect := &ParticleRect{
		rows: rows, cols: cols,
	}

	rect.System = newton.NewSystem(newton.Verlet, rows*cols)

	for i := 0; i < rect.Bodies(); i++ {

		b := rect.Body(i)

		b.SetMass(1)

		pos := rect.RestingPosition(i)

		b.Shift(pos, vect.Zero)
		b.Shift(pos, vect.Zero)

	}

	return rect
}

// The row and column in which the i-th particle is
func (rect *ParticleRect) RowAndColumn(ith int) (row, col int) {

	return i % rect.rows, i / rect.rows
}

// Initial, resting position of the i-th particle
func (rect *ParticleRect) RestingPosition(ith int) vect.Vector {

	row, col := rect.RowAndColumn(ith)

	return vect.UnitX.Scale(float64(row)).Plus(
		vect.UnitX.Scale(float64(col)),
	)
}

// Dimmensions of the rectangle
func (rect *ParticleRect) Size() (rows, cols int) {

	return rect.rows, rect.cols
}

func main() {
}
<|endoftext|>


<commit_before>// Copyright 2015, David Howden
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"os"
	"sync"

	"golang.org/x/net/websocket"

	"github.com/dhowden/httpauth"
	"github.com/dhowden/itl"

	"github.com/dhowden/tchaik/index"
	"github.com/dhowden/tchaik/store"
	"github.com/dhowden/tchaik/store/cmdflag"
)

var debug bool
var itlXML, tchLib string

var listenAddr string
var certFile, keyFile string

var auth bool

func init() {
	flag.BoolVar(&debug, "debug", false, "print debugging information")

	flag.StringVar(&listenAddr, "listen", "localhost:8080", "bind address to http listen")
	flag.StringVar(&certFile, "tls-cert", "", "path to a certificate file, must also specify -tls-key")
	flag.StringVar(&keyFile, "tls-key", "", "path to a certificate key file, must also specify -tls-cert")

	flag.StringVar(&itlXML, "itlXML", "", "path to iTunes Library XML file")
	flag.StringVar(&tchLib, "lib", "", "path to Tchaik library file")

	flag.BoolVar(&auth, "auth", false, "use basic HTTP authentication")
}

var creds = httpauth.Creds(map[string]string{
	"user": "password",
})

func readLibrary() (index.Library, error) {
	if itlXML == "" && tchLib == "" {
		return nil, fmt.Errorf("must specify one library file (-itlXML or -lib)")
	}

	if itlXML != "" && tchLib != "" {
		return nil, fmt.Errorf("must only specify one library file (-itlXML or -lib)")
	}

	var l index.Library
	if itlXML != "" {
		f, err := os.Open(itlXML)
		if err != nil {
			return nil, fmt.Errorf("could not open iTunes library file: %v", err)
		}

		fmt.Printf("Parsing %v...", itlXML)
		it, err := itl.ReadFromXML(f)
		if err != nil {
			return nil, fmt.Errorf("error parsing iTunes library file: %v", err)
		}
		f.Close()
		fmt.Println("done.")

		fmt.Printf("Building Tchaik Library...")
		l = index.Convert(index.NewITunesLibrary(&it), "TrackID")
		fmt.Println("done.")
		return l, nil
	}

	f, err := os.Open(tchLib)
	if err != nil {
		return nil, fmt.Errorf("could not open Tchaik library file: %v", err)
	}

	fmt.Printf("Parsing %v...", tchLib)
	l, err = index.ReadFrom(f)
	if err != nil {
		return nil, fmt.Errorf("error parsing Tchaik library file: %v\n", err)
	}
	fmt.Println("done.")
	return l, nil
}

func buildRootCollection(l index.Library) index.Collection {
	root := index.Collect(l, index.ByAttr(index.StringAttr("Album")))
	index.SortKeysByGroupName(root)
	return root
}

func buildSearchIndex(c index.Collection) index.Searcher {
	wi := index.BuildWordIndex(c, []string{"Composer", "Artist", "Album", "Name"})
	return index.FlatSearcher{
		Searcher: index.WordsIntersectSearcher(index.BuildPrefixExpandSearcher(wi, wi, 10)),
	}
}

func main() {
	flag.Parse()
	l, err := readLibrary()
	if err != nil {
		fmt.Printf("error: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Building root collection...")
	root := buildRootCollection(l)
	fmt.Println("done.")

	fmt.Printf("Building search index...")
	searcher := buildSearchIndex(root)
	fmt.Println("done.")

	mediaFileSystem, artworkFileSystem, err := cmdflag.Stores()
	if err != nil {
		fmt.Println("error setting up stores:", err)
		os.Exit(1)
	}

	if debug {
		mediaFileSystem = store.LogFileSystem{
			Name:      "Media",
			FileSystem: mediaFileSystem,
		}
		artworkFileSystem = store.LogFileSystem{
			Name:      "Artwork",
			FileSystem: artworkFileSystem,
		}
	}

	mediaFileSystem = &libraryFileSystem{mediaFileSystem, l}
	artworkFileSystem = &libraryFileSystem{artworkFileSystem, l}

	libAPI := LibraryAPI{
		Library:  l,
		root:     root,
		searcher: searcher,
	}

	m := buildMainHandler(libAPI, mediaFileSystem, artworkFileSystem)

	if certFile != "" && keyFile != "" {
		fmt.Printf("Web server is running on https://%v\n", listenAddr)
		fmt.Println("Quit the server with CTRL-C.")

		log.Fatal(http.ListenAndServeTLS(listenAddr, certFile, keyFile, m))
	}

	fmt.Printf("Web server is running on http://%v\n", listenAddr)
	fmt.Println("Quit the server with CTRL-C.")

	log.Fatal(http.ListenAndServe(listenAddr, m))
}

func buildMainHandler(l LibraryAPI, mediaFileSystem, artworkFileSystem http.FileSystem) http.Handler {
	var c httpauth.Checker = httpauth.None{}
	if auth {
		c = creds
	}

	w := httpauth.NewServeMux(c, http.NewServeMux())
	w.HandleFunc("/", rootHandler)
	w.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("ui/static"))))
	w.Handle("/track/", http.StripPrefix("/track/", http.FileServer(mediaFileSystem)))
	w.Handle("/artwork/", http.StripPrefix("/artwork/", http.FileServer(artworkFileSystem)))
	w.Handle("/socket", websocket.Handler(socketHandler(l)))
	return w
}

func rootHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("X-Clacks-Overhead", "GNU Terry Pratchett")
	http.ServeFile(w, r, "ui/tchaik.html")
}

func debugDumpRequest(r *http.Request) {
	if debug {
		rb, err := httputil.DumpRequest(r, true)
		if err != nil {
			fmt.Println("could not dump request:", err)
		}
		fmt.Println(string(rb))
	}
}

// Websocket handling
type socket struct {
	io.ReadWriter
	done chan struct{}
}

func (s *socket) Close() {
	select {
	case <-s.done:
		return
	default:
	}
	close(s.done)
}

type Command struct {
	Action string
	Input  string
	Path   []string
}

const (
	FetchAction  string = "FETCH"
	SearchAction string = "SEARCH"
)

func socketHandler(l LibraryAPI) func(ws *websocket.Conn) {
	return func(ws *websocket.Conn) {
		s := socket{ws, make(chan struct{})}
		out, in := make(chan interface{}), make(chan *Command)
		errCh := make(chan error)

		wg := &sync.WaitGroup{}
		wg.Add(2)

		// Encode messages from process and encode to the client
		enc := json.NewEncoder(s)
		go func() {
			defer wg.Done()
			defer s.Close()

			for {
				select {
				case x, ok := <-out:
					if !ok {
						return
					}

					if debug {
						b, err := json.MarshalIndent(x, "", "  ")
						if err != nil {
							log.Fatal(err)
						}
						fmt.Println(string(b))
					}

					if err := enc.Encode(x); err != nil {
						errCh <- fmt.Errorf("encode: %v", err)
						return
					}
				case <-s.done:
					return
				}
			}
		}()

		// Decode messages from the client and send them on the in channel
		go func() {
			defer wg.Done()
			defer s.Close()

			dec := json.NewDecoder(s)
			for {
				c := &Command{}
				if err := dec.Decode(c); err != nil {
					if err == io.EOF && debug {
						fmt.Println("websocket closed")
						return
					}
					errCh <- err
					return
				}
				in <- c
			}
		}()

		go func() {
			for x := range in {
				if debug {
					fmt.Printf("command received: %#v\n", x)
				}
				switch x.Action {
				case FetchAction:
					handleCollectionList(l, x, out)
				case SearchAction:
					handleSearch(l, x, out)
				default:
					fmt.Printf("unknown command: %v", x.Action)
				}
			}
		}()

		go func() {
			wg.Wait()

			close(in)
			close(out)
			close(errCh)
		}()

		go func() {
			for err := range errCh {
				if err != nil {
					fmt.Printf("websocket handler: %v\n", err)
				}
			}
		}()

		select {}
	}
}

func handleCollectionList(l LibraryAPI, x *Command, out chan<- interface{}) {
	if len(x.Path) < 1 {
		fmt.Printf("invalid path: %v\n", x.Path)
		return
	}

	g, err := l.Fetch(l.root, x.Path[1:])
	if err != nil {
		fmt.Printf("error in Fetch: %v (path: %#v)", err, x.Path[1:])
		return
	}

	o := struct {
		Action string
		Data   interface{}
	}{
		x.Action,
		struct {
			Path []string
			Item group
		}{
			x.Path,
			g,
		},
	}
	out <- o
}

func handleSearch(l LibraryAPI, x *Command, out chan<- interface{}) {
	paths := l.searcher.Search(x.Input)
	o := struct {
		Action string
		Data   interface{}
	}{
		Action: x.Action,
		Data:   paths,
	}
	out <- o
}
<commit_msg>Remove unused debugDumpRequest function<commit_after>// Copyright 2015, David Howden
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"sync"

	"golang.org/x/net/websocket"

	"github.com/dhowden/httpauth"
	"github.com/dhowden/itl"

	"github.com/dhowden/tchaik/index"
	"github.com/dhowden/tchaik/store"
	"github.com/dhowden/tchaik/store/cmdflag"
)

var debug bool
var itlXML, tchLib string

var listenAddr string
var certFile, keyFile string

var auth bool

func init() {
	flag.BoolVar(&debug, "debug", false, "print debugging information")

	flag.StringVar(&listenAddr, "listen", "localhost:8080", "bind address to http listen")
	flag.StringVar(&certFile, "tls-cert", "", "path to a certificate file, must also specify -tls-key")
	flag.StringVar(&keyFile, "tls-key", "", "path to a certificate key file, must also specify -tls-cert")

	flag.StringVar(&itlXML, "itlXML", "", "path to iTunes Library XML file")
	flag.StringVar(&tchLib, "lib", "", "path to Tchaik library file")

	flag.BoolVar(&auth, "auth", false, "use basic HTTP authentication")
}

var creds = httpauth.Creds(map[string]string{
	"user": "password",
})

func readLibrary() (index.Library, error) {
	if itlXML == "" && tchLib == "" {
		return nil, fmt.Errorf("must specify one library file (-itlXML or -lib)")
	}

	if itlXML != "" && tchLib != "" {
		return nil, fmt.Errorf("must only specify one library file (-itlXML or -lib)")
	}

	var l index.Library
	if itlXML != "" {
		f, err := os.Open(itlXML)
		if err != nil {
			return nil, fmt.Errorf("could not open iTunes library file: %v", err)
		}

		fmt.Printf("Parsing %v...", itlXML)
		it, err := itl.ReadFromXML(f)
		if err != nil {
			return nil, fmt.Errorf("error parsing iTunes library file: %v", err)
		}
		f.Close()
		fmt.Println("done.")

		fmt.Printf("Building Tchaik Library...")
		l = index.Convert(index.NewITunesLibrary(&it), "TrackID")
		fmt.Println("done.")
		return l, nil
	}

	f, err := os.Open(tchLib)
	if err != nil {
		return nil, fmt.Errorf("could not open Tchaik library file: %v", err)
	}

	fmt.Printf("Parsing %v...", tchLib)
	l, err = index.ReadFrom(f)
	if err != nil {
		return nil, fmt.Errorf("error parsing Tchaik library file: %v\n", err)
	}
	fmt.Println("done.")
	return l, nil
}

func buildRootCollection(l index.Library) index.Collection {
	root := index.Collect(l, index.ByAttr(index.StringAttr("Album")))
	index.SortKeysByGroupName(root)
	return root
}

func buildSearchIndex(c index.Collection) index.Searcher {
	wi := index.BuildWordIndex(c, []string{"Composer", "Artist", "Album", "Name"})
	return index.FlatSearcher{
		Searcher: index.WordsIntersectSearcher(index.BuildPrefixExpandSearcher(wi, wi, 10)),
	}
}

func main() {
	flag.Parse()
	l, err := readLibrary()
	if err != nil {
		fmt.Printf("error: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Building root collection...")
	root := buildRootCollection(l)
	fmt.Println("done.")

	fmt.Printf("Building search index...")
	searcher := buildSearchIndex(root)
	fmt.Println("done.")

	mediaFileSystem, artworkFileSystem, err := cmdflag.Stores()
	if err != nil {
		fmt.Println("error setting up stores:", err)
		os.Exit(1)
	}

	if debug {
		mediaFileSystem = store.LogFileSystem{
			Name:      "Media",
			FileSystem: mediaFileSystem,
		}
		artworkFileSystem = store.LogFileSystem{
			Name:      "Artwork",
			FileSystem: artworkFileSystem,
		}
	}

	mediaFileSystem = &libraryFileSystem{mediaFileSystem, l}
	artworkFileSystem = &libraryFileSystem{artworkFileSystem, l}

	libAPI := LibraryAPI{
		Library:  l,
		root:     root,
		searcher: searcher,
	}

	m := buildMainHandler(libAPI, mediaFileSystem, artworkFileSystem)

	if certFile != "" && keyFile != "" {
		fmt.Printf("Web server is running on https://%v\n", listenAddr)
		fmt.Println("Quit the server with CTRL-C.")

		log.Fatal(http.ListenAndServeTLS(listenAddr, certFile, keyFile, m))
	}

	fmt.Printf("Web server is running on http://%v\n", listenAddr)
	fmt.Println("Quit the server with CTRL-C.")

	log.Fatal(http.ListenAndServe(listenAddr, m))
}

func buildMainHandler(l LibraryAPI, mediaFileSystem, artworkFileSystem http.FileSystem) http.Handler {
	var c httpauth.Checker = httpauth.None{}
	if auth {
		c = creds
	}

	w := httpauth.NewServeMux(c, http.NewServeMux())
	w.HandleFunc("/", rootHandler)
	w.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("ui/static"))))
	w.Handle("/track/", http.StripPrefix("/track/", http.FileServer(mediaFileSystem)))
	w.Handle("/artwork/", http.StripPrefix("/artwork/", http.FileServer(artworkFileSystem)))
	w.Handle("/socket", websocket.Handler(socketHandler(l)))
	return w
}

func rootHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("X-Clacks-Overhead", "GNU Terry Pratchett")
	http.ServeFile(w, r, "ui/tchaik.html")
}

// Websocket handling
type socket struct {
	io.ReadWriter
	done chan struct{}
}

func (s *socket) Close() {
	select {
	case <-s.done:
		return
	default:
	}
	close(s.done)
}

type Command struct {
	Action string
	Input  string
	Path   []string
}

const (
	FetchAction  string = "FETCH"
	SearchAction string = "SEARCH"
)

func socketHandler(l LibraryAPI) func(ws *websocket.Conn) {
	return func(ws *websocket.Conn) {
		s := socket{ws, make(chan struct{})}
		out, in := make(chan interface{}), make(chan *Command)
		errCh := make(chan error)

		wg := &sync.WaitGroup{}
		wg.Add(2)

		// Encode messages from process and encode to the client
		enc := json.NewEncoder(s)
		go func() {
			defer wg.Done()
			defer s.Close()

			for {
				select {
				case x, ok := <-out:
					if !ok {
						return
					}

					if debug {
						b, err := json.MarshalIndent(x, "", "  ")
						if err != nil {
							log.Fatal(err)
						}
						fmt.Println(string(b))
					}

					if err := enc.Encode(x); err != nil {
						errCh <- fmt.Errorf("encode: %v", err)
						return
					}
				case <-s.done:
					return
				}
			}
		}()

		// Decode messages from the client and send them on the in channel
		go func() {
			defer wg.Done()
			defer s.Close()

			dec := json.NewDecoder(s)
			for {
				c := &Command{}
				if err := dec.Decode(c); err != nil {
					if err == io.EOF && debug {
						fmt.Println("websocket closed")
						return
					}
					errCh <- err
					return
				}
				in <- c
			}
		}()

		go func() {
			for x := range in {
				if debug {
					fmt.Printf("command received: %#v\n", x)
				}
				switch x.Action {
				case FetchAction:
					handleCollectionList(l, x, out)
				case SearchAction:
					handleSearch(l, x, out)
				default:
					fmt.Printf("unknown command: %v", x.Action)
				}
			}
		}()

		go func() {
			wg.Wait()

			close(in)
			close(out)
			close(errCh)
		}()

		go func() {
			for err := range errCh {
				if err != nil {
					fmt.Printf("websocket handler: %v\n", err)
				}
			}
		}()

		select {}
	}
}

func handleCollectionList(l LibraryAPI, x *Command, out chan<- interface{}) {
	if len(x.Path) < 1 {
		fmt.Printf("invalid path: %v\n", x.Path)
		return
	}

	g, err := l.Fetch(l.root, x.Path[1:])
	if err != nil {
		fmt.Printf("error in Fetch: %v (path: %#v)", err, x.Path[1:])
		return
	}

	o := struct {
		Action string
		Data   interface{}
	}{
		x.Action,
		struct {
			Path []string
			Item group
		}{
			x.Path,
			g,
		},
	}
	out <- o
}

func handleSearch(l LibraryAPI, x *Command, out chan<- interface{}) {
	paths := l.searcher.Search(x.Input)
	o := struct {
		Action string
		Data   interface{}
	}{
		Action: x.Action,
		Data:   paths,
	}
	out <- o
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"github.com/whosonfirst/go-whosonfirst-index"
	"github.com/whosonfirst/go-whosonfirst-s3/sync"
	"log"
	"runtime"
	"sync/atomic"
	"time"
)

func main() {

	var mode = flag.String("mode", "files", "...")
	var region = flag.String("region", "us-east-1", "...")
	var bucket = flag.String("bucket", "whosonfirst.mapzen.com", "...")
	var prefix = flag.String("prefix", "", "...")
	var acl = flag.String("acl", "public-read", "...")
	var creds = flag.String("credentials", "default", "...")
	var ratelimit = flag.Int("rate-limit", 100000, "...")
	var dryrun = flag.Bool("dryrun", false, "...")
	var force = flag.Bool("force", false, "...")
	var verbose = flag.Bool("verbose", false, "...")
	var procs = flag.Int("processes", (runtime.NumCPU() * 2), "The number of concurrent processes to clone data with")

	flag.Parse()

	runtime.GOMAXPROCS(*procs)

	opts := sync.RemoteSyncOptions{
		Region:      *region,
		Bucket:      *bucket,
		Prefix:      *prefix,
		ACL:         *acl,
		RateLimit:   *ratelimit,
		Dryrun:      *dryrun,
		Force:       *force,
		Verbose:     *verbose,
		Credentials: *creds,
	}

	sync, err := sync.NewRemoteSync(opts)

	if err != nil {
		log.Fatal(err)
	}

	cb, err := sync.SyncFunc()

	if err != nil {
		log.Fatal(err)
	}

	idx, err := index.NewIndexer(*mode, cb)

	if err != nil {
		log.Fatal(err)
	}

	done_ch := make(chan bool)

	go func() {

		for {
			select {
			case <-done_ch:
				break
			case <-time.After(1 * time.Minute):
				i := atomic.LoadInt64(&idx.Indexed) // please just make this part of go-whosonfirst-index
				log.Printf("%d indexed\n", i)
			}
		}
	}()

	t1 := time.Now()

	for _, path := range flag.Args() {

		ta := time.Now()

		err := idx.IndexPath(path)

		if err != nil {
			log.Println(err)
			break
		}

		tb := time.Since(ta)
		log.Printf("time to index %v\n", tb)
	}

	// this code doesn't exist and I am not sure how I want to deal
	// with it yet (20171212/thisisaaronland)

	/*

		if sync.HasRetries() {

			idx, err = index.NewIndexer("files", cb)

			if err != nil {
				log.Fatal(err)
			}

			for _, path := range sync.Retries() {

				idx.IndexPath(path)
			}
		}

	*/

	done_ch <- true

	t2 := time.Since(t1)
	i := atomic.LoadInt64(&idx.Indexed) // see above

	log.Printf("time to index %d documents : %v\n", i, t2)
}
<commit_msg>better reporting<commit_after>package main

import (
	"flag"
	"github.com/whosonfirst/go-whosonfirst-index"
	"github.com/whosonfirst/go-whosonfirst-s3/sync"
	"log"
	"runtime"
	"sync/atomic"
	"time"
)

func main() {

	var mode = flag.String("mode", "files", "...")
	var region = flag.String("region", "us-east-1", "...")
	var bucket = flag.String("bucket", "whosonfirst.mapzen.com", "...")
	var prefix = flag.String("prefix", "", "...")
	var acl = flag.String("acl", "public-read", "...")
	var creds = flag.String("credentials", "default", "...")
	var ratelimit = flag.Int("rate-limit", 100000, "...")
	var dryrun = flag.Bool("dryrun", false, "...")
	var force = flag.Bool("force", false, "...")
	var verbose = flag.Bool("verbose", false, "...")
	var procs = flag.Int("processes", (runtime.NumCPU() * 2), "The number of concurrent processes to clone data with")

	flag.Parse()

	runtime.GOMAXPROCS(*procs)

	opts := sync.RemoteSyncOptions{
		Region:      *region,
		Bucket:      *bucket,
		Prefix:      *prefix,
		ACL:         *acl,
		RateLimit:   *ratelimit,
		Dryrun:      *dryrun,
		Force:       *force,
		Verbose:     *verbose,
		Credentials: *creds,
	}

	sync, err := sync.NewRemoteSync(opts)

	if err != nil {
		log.Fatal(err)
	}

	cb, err := sync.SyncFunc()

	if err != nil {
		log.Fatal(err)
	}

	idx, err := index.NewIndexer(*mode, cb)

	if err != nil {
		log.Fatal(err)
	}

	done_ch := make(chan bool)

	go func() {

		for {
			select {
			case <-done_ch:
				break
			case <-time.After(1 * time.Minute):
				i := atomic.LoadInt64(&idx.Indexed) // please just make this part of go-whosonfirst-index
				log.Printf("%d indexed\n", i)
			}
		}
	}()

	t1 := time.Now()

	for _, path := range flag.Args() {

		ta := time.Now()

		err := idx.IndexPath(path)

		if err != nil {
			log.Println(err)
			break
		}

		tb := time.Since(ta)
		log.Printf("time to index %s : %v\n", path, tb)
	}

	// this code doesn't exist and I am not sure how I want to deal
	// with it yet (20171212/thisisaaronland)

	/*

		if sync.HasRetries() {

			idx, err = index.NewIndexer("files", cb)

			if err != nil {
				log.Fatal(err)
			}

			for _, path := range sync.Retries() {

				idx.IndexPath(path)
			}
		}

	*/

	done_ch <- true

	t2 := time.Since(t1)
	i := atomic.LoadInt64(&idx.Indexed) // see above

	log.Printf("time to index %d documents : %v\n", i, t2)
}
<|endoftext|>


<commit_before>// Copyright 2016 Sorint.lab
//
// 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 integration

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
	"time"

	"github.com/satori/go.uuid"
	"github.com/sorintlab/stolon/common"
	"github.com/sorintlab/stolon/pkg/cluster"
	"github.com/sorintlab/stolon/pkg/store"
)

func TestPITR(t *testing.T) {
	t.Parallel()

	dir, err := ioutil.TempDir("", "stolon")
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	defer os.RemoveAll(dir)

	baseBackupDir, err := ioutil.TempDir(dir, "basebackup")
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	archiveBackupDir, err := ioutil.TempDir(dir, "archivebackup")
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	tstore := setupStore(t, dir)
	defer tstore.Stop()

	storeEndpoints := fmt.Sprintf("%s:%s", tstore.listenAddress, tstore.port)

	clusterName := uuid.NewV4().String()

	storePath := filepath.Join(common.StoreBasePath, clusterName)

	sm := store.NewStoreManager(tstore.store, storePath)

	initialClusterSpec := &cluster.ClusterSpec{
		InitMode:           cluster.ClusterInitModeP(cluster.ClusterInitModeNew),
		SleepInterval:      &cluster.Duration{Duration: 2 * time.Second},
		FailInterval:       &cluster.Duration{Duration: 5 * time.Second},
		ConvergenceTimeout: &cluster.Duration{Duration: 30 * time.Second},
		PGParameters: cluster.PGParameters{
			"archive_mode":    "on",
			"archive_command": fmt.Sprintf("cp %%p %s/%%f", archiveBackupDir),
		},
	}
	initialClusterSpecFile, err := writeClusterSpec(dir, initialClusterSpec)
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	tk, err := NewTestKeeper(t, dir, clusterName, pgSUUsername, pgSUPassword, pgReplUsername, pgReplPassword, tstore.storeBackend, storeEndpoints)
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := tk.Start(); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	defer tk.Stop()

	ts, err := NewTestSentinel(t, dir, clusterName, tstore.storeBackend, storeEndpoints, fmt.Sprintf("--initial-cluster-spec=%s", initialClusterSpecFile))
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := ts.Start(); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	// Wait for clusterView containing a master
	_, err = WaitClusterDataWithMaster(sm, 30*time.Second)
	if err != nil {
		t.Fatal("expected a master in cluster view")
	}
	if err := tk.WaitDBUp(60 * time.Second); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := tk.WaitRole(common.RoleMaster, 30*time.Second); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := populate(t, tk); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := write(t, tk, 2, 2); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	// ioutil.Tempfile already creates files with 0600 permissions
	pgpass, err := ioutil.TempFile("", "pgpass")
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	pgpass.WriteString(fmt.Sprintf("%s:%s:*:%s:%s\n", tk.pgListenAddress, tk.pgPort, tk.pgReplUsername, tk.pgReplPassword))
	// Don't save the wal during the basebackup (-x). This to test that archive_command and restore command correctly work.
	cmd := exec.Command("pg_basebackup", "-F", "tar", "-D", baseBackupDir, "-h", tk.pgListenAddress, "-p", tk.pgPort, "-U", tk.pgReplUsername)
	cmd.Env = append(cmd.Env, fmt.Sprintf("PGPASSFILE=%s", pgpass.Name()))
	t.Logf("execing cmd: %s", cmd)
	if out, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("error: %v, output: %s", err, string(out))
	}

	// Switch wal so they will be archived
	if _, err := tk.db.Exec("select pg_switch_xlog()"); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	ts.Stop()

	// Delete the current cluster data
	if err := tstore.store.Delete(filepath.Join(storePath, "clusterdata")); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	// Delete sentinel leader key to just speedup new election
	if err := tstore.store.Delete(filepath.Join(storePath, common.SentinelLeaderKey)); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	// Now initialize a new cluster with the existing keeper
	initialClusterSpec = &cluster.ClusterSpec{
		InitMode:           cluster.ClusterInitModeP(cluster.ClusterInitModePITR),
		SleepInterval:      &cluster.Duration{Duration: 2 * time.Second},
		FailInterval:       &cluster.Duration{Duration: 5 * time.Second},
		ConvergenceTimeout: &cluster.Duration{Duration: 30 * time.Second},
		PITRConfig: &cluster.PITRConfig{
			DataRestoreCommand: fmt.Sprintf("tar xvf %s/base.tar -C %%d", baseBackupDir),
			ArchiveRecoverySettings: &cluster.ArchiveRecoverySettings{
				RestoreCommand: fmt.Sprintf("cp %s/%%f %%p", archiveBackupDir),
			},
		},
	}
	initialClusterSpecFile, err = writeClusterSpec(dir, initialClusterSpec)
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	ts, err = NewTestSentinel(t, dir, clusterName, tstore.storeBackend, storeEndpoints, fmt.Sprintf("--initial-cluster-spec=%s", initialClusterSpecFile))
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := ts.Start(); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	defer ts.Stop()

	if err := WaitClusterPhase(sm, cluster.ClusterPhaseNormal, 60*time.Second); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	_, err = WaitClusterDataWithMaster(sm, 30*time.Second)
	if err != nil {
		t.Fatal("expected a master in cluster view")
	}
	if err := tk.WaitDBUp(60 * time.Second); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := tk.WaitRole(common.RoleMaster, 30*time.Second); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	if err := tk.WaitDBUp(60 * time.Second); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	c, err := getLines(t, tk)
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if c != 1 {
		t.Fatalf("wrong number of lines, want: %d, got: %d", 2, c)
	}

}
<commit_msg>pitr tests: check init spec pgParameters are honored<commit_after>// Copyright 2016 Sorint.lab
//
// 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 integration

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
	"time"

	"github.com/satori/go.uuid"
	"github.com/sorintlab/stolon/common"
	"github.com/sorintlab/stolon/pkg/cluster"
	"github.com/sorintlab/stolon/pkg/store"
)

func TestPITR(t *testing.T) {
	t.Parallel()

	dir, err := ioutil.TempDir("", "stolon")
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	defer os.RemoveAll(dir)

	baseBackupDir, err := ioutil.TempDir(dir, "basebackup")
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	archiveBackupDir, err := ioutil.TempDir(dir, "archivebackup")
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	tstore := setupStore(t, dir)
	defer tstore.Stop()

	storeEndpoints := fmt.Sprintf("%s:%s", tstore.listenAddress, tstore.port)

	clusterName := uuid.NewV4().String()

	storePath := filepath.Join(common.StoreBasePath, clusterName)

	sm := store.NewStoreManager(tstore.store, storePath)

	initialClusterSpec := &cluster.ClusterSpec{
		InitMode:           cluster.ClusterInitModeP(cluster.ClusterInitModeNew),
		SleepInterval:      &cluster.Duration{Duration: 2 * time.Second},
		FailInterval:       &cluster.Duration{Duration: 5 * time.Second},
		ConvergenceTimeout: &cluster.Duration{Duration: 30 * time.Second},
		PGParameters: cluster.PGParameters{
			"archive_mode":    "on",
			"archive_command": fmt.Sprintf("cp %%p %s/%%f", archiveBackupDir),
		},
	}
	initialClusterSpecFile, err := writeClusterSpec(dir, initialClusterSpec)
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	tk, err := NewTestKeeper(t, dir, clusterName, pgSUUsername, pgSUPassword, pgReplUsername, pgReplPassword, tstore.storeBackend, storeEndpoints)
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := tk.Start(); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	defer tk.Stop()

	ts, err := NewTestSentinel(t, dir, clusterName, tstore.storeBackend, storeEndpoints, fmt.Sprintf("--initial-cluster-spec=%s", initialClusterSpecFile))
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := ts.Start(); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	// Wait for clusterView containing a master
	_, err = WaitClusterDataWithMaster(sm, 30*time.Second)
	if err != nil {
		t.Fatal("expected a master in cluster view")
	}
	if err := tk.WaitDBUp(60 * time.Second); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := tk.WaitRole(common.RoleMaster, 30*time.Second); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := populate(t, tk); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := write(t, tk, 2, 2); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	// ioutil.Tempfile already creates files with 0600 permissions
	pgpass, err := ioutil.TempFile("", "pgpass")
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	pgpass.WriteString(fmt.Sprintf("%s:%s:*:%s:%s\n", tk.pgListenAddress, tk.pgPort, tk.pgReplUsername, tk.pgReplPassword))
	// Don't save the wal during the basebackup (-x). This to test that archive_command and restore command correctly work.
	cmd := exec.Command("pg_basebackup", "-F", "tar", "-D", baseBackupDir, "-h", tk.pgListenAddress, "-p", tk.pgPort, "-U", tk.pgReplUsername)
	cmd.Env = append(cmd.Env, fmt.Sprintf("PGPASSFILE=%s", pgpass.Name()))
	t.Logf("execing cmd: %v", cmd)
	if out, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("error: %v, output: %s", err, string(out))
	}

	// Switch wal so they will be archived
	if _, err := tk.db.Exec("select pg_switch_xlog()"); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	ts.Stop()

	// Delete the current cluster data
	if err := tstore.store.Delete(filepath.Join(storePath, "clusterdata")); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	// Delete sentinel leader key to just speedup new election
	if err := tstore.store.Delete(filepath.Join(storePath, common.SentinelLeaderKey)); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	// Now initialize a new cluster with the existing keeper
	initialClusterSpec = &cluster.ClusterSpec{
		InitMode:           cluster.ClusterInitModeP(cluster.ClusterInitModePITR),
		SleepInterval:      &cluster.Duration{Duration: 2 * time.Second},
		FailInterval:       &cluster.Duration{Duration: 5 * time.Second},
		ConvergenceTimeout: &cluster.Duration{Duration: 30 * time.Second},
		PITRConfig: &cluster.PITRConfig{
			DataRestoreCommand: fmt.Sprintf("tar xvf %s/base.tar -C %%d", baseBackupDir),
			ArchiveRecoverySettings: &cluster.ArchiveRecoverySettings{
				RestoreCommand: fmt.Sprintf("cp %s/%%f %%p", archiveBackupDir),
			},
		},
		PGParameters: cluster.PGParameters{
			"max_prepared_transactions": "100",
		},
	}
	initialClusterSpecFile, err = writeClusterSpec(dir, initialClusterSpec)
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	ts, err = NewTestSentinel(t, dir, clusterName, tstore.storeBackend, storeEndpoints, fmt.Sprintf("--initial-cluster-spec=%s", initialClusterSpecFile))
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := ts.Start(); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	defer ts.Stop()

	if err := WaitClusterPhase(sm, cluster.ClusterPhaseNormal, 60*time.Second); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	_, err = WaitClusterDataWithMaster(sm, 30*time.Second)
	if err != nil {
		t.Fatal("expected a master in cluster view")
	}
	if err := tk.WaitDBUp(60 * time.Second); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if err := tk.WaitRole(common.RoleMaster, 30*time.Second); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	if err := tk.WaitDBUp(60 * time.Second); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	c, err := getLines(t, tk)
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if c != 1 {
		t.Fatalf("wrong number of lines, want: %d, got: %d", 2, c)
	}

	pgParameters, err := tk.GetPGParameters()
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}
	if v := pgParameters["max_prepared_transactions"]; v != "100" {
		t.Fatalf("expected max_prepared_transactions == 100 got %q", v)
	}
}
<|endoftext|>