<commit_before>//
// Site-checker runs a slice of commands on a slice of Drush aliases matching a specified pattern.
//
package main

import (
	"flag"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/fubarhouse/golang-drush/command"
	"os"
	"strings"
)

func main() {
	var strAliases = flag.String("aliases", "", "comma-separated list of aliases for action")
	var strCommands = flag.String("commands", "rr,updb --yes,cc all", "comma-separated list of commands for action")
	var strPattern = flag.String("pattern", "%v", "A modifier which allows rewriting of aliases replacing '%v' in the pattern with the alias.")
	var boolVerbose = flag.Bool("verbose", false, "adds raw output to end of program.")
	flag.Parse()

	// Remove double spaces.
	*strAliases = strings.Replace(*strAliases, "  ", " ", -1)
	*strCommands = strings.Replace(*strCommands, "  ", " ", -1)

	var FinalOutput []string

	if !strings.Contains(*strPattern, "%v") {
		log.Errorln("Specified pattern does not include alias modifier.")
	}

	if *strCommands == "" {
		log.Errorln("Commands are not specified.")
	}

	if *strAliases == "" {
		flag.Usage()
		log.Errorln("Aliases are not specified.")
	} else {
		*strAliases = strings.Replace(*strAliases, " ", "", -1)
	}

	if *strAliases != "" && *strCommands != "" {
		for _, Alias := range strings.Split(*strAliases, ",") {
			Alias = strings.Replace(*strPattern, "%v", Alias, 1)
			Alias = strings.Trim(Alias, " ")
			for _, Command := range strings.Split(*strCommands, ",") {
				Command = strings.Trim(Command, " ")
				FinalOutput = append(FinalOutput, fmt.Sprintf("drush @%v '%v'\n", Alias, Command))
				DrushCommand := command.NewDrushCommand()
				DrushCommand.SetAlias(Alias)
				DrushCommand.SetCommand(Command)
				if Command != "" {
					DrushCommandOut, DrushCommandError := DrushCommand.Output()
					if DrushCommandError != nil {
						log.Warnf("drush %v '%v' was unsuccessful.", DrushCommand.GetAlias(), DrushCommand.GetCommand())
						if *boolVerbose {
							StdOutLines := DrushCommandOut
							for _, StdOutLine := range StdOutLines {
								FinalOutput = append(FinalOutput, StdOutLine)
							}
						}
					} else {
						log.Infof("drush %v '%v' was successful.", DrushCommand.GetAlias(), DrushCommand.GetCommand())
						if *boolVerbose {
							StdOutLines := DrushCommandOut
							for _, StdOutLine := range StdOutLines {
								FinalOutput = append(FinalOutput, StdOutLine)
							}
						}
					}
				}
			}
		}
		if *boolVerbose {
			for _, value := range FinalOutput {
				value = strings.Replace(value, "\n", "", -1)
				value = strings.Replace(value, "\t", "", -1)
				value = strings.Replace(value, "\r", "", -1)
				value = fmt.Sprint(value)
				if value == "Permission denied (publickey)." {
					log.Warnln(value)
				} else {
					log.Infoln(value)
				}
			}
		}
	} else {
		flag.Usage()
		os.Exit(1)
	}
}
<commit_msg>Do not replace new lines on verbose output to retain clean formatting of verbose output.<commit_after>//
// Site-checker runs a slice of commands on a slice of Drush aliases matching a specified pattern.
//
package main

import (
	"flag"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/fubarhouse/golang-drush/command"
	"os"
	"strings"
)

func main() {
	var strAliases = flag.String("aliases", "", "comma-separated list of aliases for action")
	var strCommands = flag.String("commands", "rr,updb --yes,cc all", "comma-separated list of commands for action")
	var strPattern = flag.String("pattern", "%v", "A modifier which allows rewriting of aliases replacing '%v' in the pattern with the alias.")
	var boolVerbose = flag.Bool("verbose", false, "adds raw output to end of program.")
	flag.Parse()

	// Remove double spaces.
	*strAliases = strings.Replace(*strAliases, "  ", " ", -1)
	*strCommands = strings.Replace(*strCommands, "  ", " ", -1)

	var FinalOutput []string

	if !strings.Contains(*strPattern, "%v") {
		log.Errorln("Specified pattern does not include alias modifier.")
	}

	if *strCommands == "" {
		log.Errorln("Commands are not specified.")
	}

	if *strAliases == "" {
		flag.Usage()
		log.Errorln("Aliases are not specified.")
	} else {
		*strAliases = strings.Replace(*strAliases, " ", "", -1)
	}

	if *strAliases != "" && *strCommands != "" {
		for _, Alias := range strings.Split(*strAliases, ",") {
			Alias = strings.Replace(*strPattern, "%v", Alias, 1)
			Alias = strings.Trim(Alias, " ")
			for _, Command := range strings.Split(*strCommands, ",") {
				Command = strings.Trim(Command, " ")
				FinalOutput = append(FinalOutput, fmt.Sprintf("drush @%v '%v'\n", Alias, Command))
				DrushCommand := command.NewDrushCommand()
				DrushCommand.SetAlias(Alias)
				DrushCommand.SetCommand(Command)
				if Command != "" {
					DrushCommandOut, DrushCommandError := DrushCommand.Output()
					if DrushCommandError != nil {
						log.Warnf("drush %v '%v' was unsuccessful.", DrushCommand.GetAlias(), DrushCommand.GetCommand())
						if *boolVerbose {
							StdOutLines := DrushCommandOut
							for _, StdOutLine := range StdOutLines {
								FinalOutput = append(FinalOutput, StdOutLine)
							}
						}
					} else {
						log.Infof("drush %v '%v' was successful.", DrushCommand.GetAlias(), DrushCommand.GetCommand())
						if *boolVerbose {
							StdOutLines := DrushCommandOut
							for _, StdOutLine := range StdOutLines {
								FinalOutput = append(FinalOutput, StdOutLine)
							}
						}
					}
				}
			}
		}
		if *boolVerbose {
			for _, value := range FinalOutput {
				value = strings.Replace(value, "\t", "", -1)
				value = strings.Replace(value, "\r", "", -1)
				value = fmt.Sprint(value)
				if value == "Permission denied (publickey)." {
					log.Warnln(value)
				} else {
					log.Infoln(value)
				}
			}
		}
	} else {
		flag.Usage()
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015-2016 trivago GmbH
//
// 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 contrib

// This is a stub file to enable registration of vendor specific plugins that
// are placed in sub folders of this folder.

import (
	_ "github.com/trivago/gollum/contrib/native" // plugins using cgo native bindings
	//_ "your/company/package"
)

func init() {
}
<commit_msg>disabled native plugin loading<commit_after>// Copyright 2015-2016 trivago GmbH
//
// 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 contrib

// This is a stub file to enable registration of vendor specific plugins that
// are placed in sub folders of this folder.

import (
//_ "github.com/trivago/gollum/contrib/native" // plugins using cgo native bindings
//_ "your/company/package"
)

func init() {
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"os/signal"
	"path"
	"sync"
	"time"

	"github.com/elazarl/goproxy"
	"github.com/elazarl/goproxy/transport"
)

type FileStream struct {
	path string
	f    *os.File
}

func NewFileStream(path string) *FileStream {
	return &FileStream{path, nil}
}

func (fs *FileStream) Write(b []byte) (nr int, err error) {
	if fs.f == nil {
		fs.f, err = os.Create(fs.path)
		if err != nil {
			return 0, err
		}
	}
	return fs.f.Write(b)
}

func (fs *FileStream) Close() error {
	fmt.Println("Close", fs.path)
	if fs.f == nil {
		return errors.New("FileStream was never written into")
	}
	return fs.f.Close()
}

type Meta struct {
	req      *http.Request
	resp     *http.Response
	err      error
	t        time.Time
	sess     int64
	bodyPath string
	from     string
}

func fprintf(nr *int64, err *error, w io.Writer, pat string, a ...interface{}) {
	if *err != nil {
		return
	}
	var n int
	n, *err = fmt.Fprintf(w, pat, a...)
	*nr += int64(n)
}

func write(nr *int64, err *error, w io.Writer, b []byte) {
	if *err != nil {
		return
	}
	var n int
	n, *err = w.Write(b)
	*nr += int64(n)
}

func (m *Meta) WriteTo(w io.Writer) (nr int64, err error) {
	if m.req != nil {
		fprintf(&nr, &err, w, "Type: request\r\n")
	} else if m.resp != nil {
		fprintf(&nr, &err, w, "Type: response\r\n")
	}
	fprintf(&nr, &err, w, "ReceivedAt: %v\r\n", m.t)
	fprintf(&nr, &err, w, "Session: %d\r\n", m.sess)
	fprintf(&nr, &err, w, "From: %v\r\n", m.from)
	if m.err != nil {
		// note the empty response
		fprintf(&nr, &err, w, "Error: %v\r\n\r\n\r\n\r\n", m.err)
	} else if m.req != nil {
		fprintf(&nr, &err, w, "\r\n")
		buf, err2 := httputil.DumpRequest(m.req, false)
		if err2 != nil {
			return nr, err2
		}
		write(&nr, &err, w, buf)
	} else if m.resp != nil {
		fprintf(&nr, &err, w, "\r\n")
		buf, err2 := httputil.DumpResponse(m.resp, false)
		if err2 != nil {
			return nr, err2
		}
		write(&nr, &err, w, buf)
	}
	return
}

// HttpLogger is an asynchronous HTTP request/response logger. It traces
// requests and responses headers in a "log" file in logger directory and dumps
// their bodies in files prefixed with the session identifiers.
// Close it to ensure pending items are correctly logged.
type HttpLogger struct {
	path  string
	c     chan *Meta
	errch chan error
}

func NewLogger(basepath string) (*HttpLogger, error) {
	f, err := os.Create(path.Join(basepath, "log"))
	if err != nil {
		return nil, err
	}
	logger := &HttpLogger{basepath, make(chan *Meta), make(chan error)}
	go func() {
		for m := range logger.c {
			if _, err := m.WriteTo(f); err != nil {
				log.Println("Can't write meta", err)
			}
		}
		logger.errch <- f.Close()
	}()
	return logger, nil
}

func (logger *HttpLogger) LogResp(resp *http.Response, ctx *goproxy.ProxyCtx) {
	body := path.Join(logger.path, fmt.Sprintf("%d_resp", ctx.Session))
	from := ""
	if ctx.UserData != nil {
		from = ctx.UserData.(*transport.RoundTripDetails).TCPAddr.String()
	}
	if resp == nil {
		resp = emptyResp
	} else {
		resp.Body = NewTeeReadCloser(resp.Body, NewFileStream(body))
	}
	logger.LogMeta(&Meta{
		resp: resp,
		err:  ctx.Error,
		t:    time.Now(),
		sess: ctx.Session,
		from: from})
}

var emptyResp = &http.Response{}
var emptyReq = &http.Request{}

func (logger *HttpLogger) LogReq(req *http.Request, ctx *goproxy.ProxyCtx) {
	body := path.Join(logger.path, fmt.Sprintf("%d_req", ctx.Session))
	if req == nil {
		req = emptyReq
	} else {
		req.Body = NewTeeReadCloser(req.Body, NewFileStream(body))
	}
	logger.LogMeta(&Meta{
		req:  req,
		err:  ctx.Error,
		t:    time.Now(),
		sess: ctx.Session,
		from: req.RemoteAddr})
}

func (logger *HttpLogger) LogMeta(m *Meta) {
	logger.c <- m
}

func (logger *HttpLogger) Close() error {
	close(logger.c)
	return <-logger.errch
}

// TeeReadCloser extends io.TeeReader by allowing reader and writer to be
// closed.
type TeeReadCloser struct {
	r io.Reader
	w io.WriteCloser
	c io.Closer
}

func NewTeeReadCloser(r io.ReadCloser, w io.WriteCloser) io.ReadCloser {
	return &TeeReadCloser{io.TeeReader(r, w), w, r}
}

func (t *TeeReadCloser) Read(b []byte) (int, error) {
	return t.r.Read(b)
}

// Close attempts to close the reader and write. It returns an error if both
// failed to Close.
func (t *TeeReadCloser) Close() error {
	err1 := t.c.Close()
	err2 := t.w.Close()
	if err1 != nil {
		return err2
	}
	return err1
}

// stoppableListener serves stoppableConn and tracks their lifetime to notify
// when it is safe to terminate the application.
type stoppableListener struct {
	net.Listener
	sync.WaitGroup
}

type stoppableConn struct {
	net.Conn
	wg *sync.WaitGroup
}

func newStoppableListener(l net.Listener) *stoppableListener {
	return &stoppableListener{l, sync.WaitGroup{}}
}

func (sl *stoppableListener) Accept() (net.Conn, error) {
	c, err := sl.Listener.Accept()
	if err != nil {
		return c, err
	}
	sl.Add(1)
	return &stoppableConn{c, &sl.WaitGroup}, nil
}

func (sc *stoppableConn) Close() error {
	sc.wg.Done()
	return sc.Conn.Close()
}

func main() {
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	addr := flag.String("l", ":8080", "on which address should the proxy listen")
	flag.Parse()
	proxy := goproxy.NewProxyHttpServer()
	proxy.Verbose = *verbose
	if err := os.MkdirAll("db", 0755); err != nil {
		log.Fatal("Can't create dir", err)
	}
	logger, err := NewLogger("db")
	if err != nil {
		log.Fatal("can't open log file", err)
	}
	tr := transport.Transport{Proxy: transport.ProxyFromEnvironment}
	// For every incoming request, override the RoundTripper to extract
	// connection information. Store it is session context log it after
	// handling the response.
	proxy.OnRequest().DoFunc(func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		ctx.RoundTripper = goproxy.RoundTripperFunc(func(req *http.Request, ctx *goproxy.ProxyCtx) (resp *http.Response, err error) {
			ctx.UserData, resp, err = tr.DetailedRoundTrip(req)
			return
		})
		logger.LogReq(req, ctx)
		return req, nil
	})
	proxy.OnResponse().DoFunc(func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
		logger.LogResp(resp, ctx)
		return resp
	})
	l, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatal("listen:", err)
	}
	sl := newStoppableListener(l)
	ch := make(chan os.Signal)
	signal.Notify(ch, os.Interrupt)
	go func() {
		<-ch
		log.Println("Got SIGINT exiting")
		sl.Add(1)
		sl.Close()
		logger.Close()
		sl.Done()
	}()
	log.Println("Starting Proxy")
	http.Serve(sl, proxy)
	sl.Wait()
	log.Println("All connections closed - exit")
}
<commit_msg>httpdump: make TeeReadCloser.Close fail if either output fails<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"os/signal"
	"path"
	"sync"
	"time"

	"github.com/elazarl/goproxy"
	"github.com/elazarl/goproxy/transport"
)

type FileStream struct {
	path string
	f    *os.File
}

func NewFileStream(path string) *FileStream {
	return &FileStream{path, nil}
}

func (fs *FileStream) Write(b []byte) (nr int, err error) {
	if fs.f == nil {
		fs.f, err = os.Create(fs.path)
		if err != nil {
			return 0, err
		}
	}
	return fs.f.Write(b)
}

func (fs *FileStream) Close() error {
	fmt.Println("Close", fs.path)
	if fs.f == nil {
		return errors.New("FileStream was never written into")
	}
	return fs.f.Close()
}

type Meta struct {
	req      *http.Request
	resp     *http.Response
	err      error
	t        time.Time
	sess     int64
	bodyPath string
	from     string
}

func fprintf(nr *int64, err *error, w io.Writer, pat string, a ...interface{}) {
	if *err != nil {
		return
	}
	var n int
	n, *err = fmt.Fprintf(w, pat, a...)
	*nr += int64(n)
}

func write(nr *int64, err *error, w io.Writer, b []byte) {
	if *err != nil {
		return
	}
	var n int
	n, *err = w.Write(b)
	*nr += int64(n)
}

func (m *Meta) WriteTo(w io.Writer) (nr int64, err error) {
	if m.req != nil {
		fprintf(&nr, &err, w, "Type: request\r\n")
	} else if m.resp != nil {
		fprintf(&nr, &err, w, "Type: response\r\n")
	}
	fprintf(&nr, &err, w, "ReceivedAt: %v\r\n", m.t)
	fprintf(&nr, &err, w, "Session: %d\r\n", m.sess)
	fprintf(&nr, &err, w, "From: %v\r\n", m.from)
	if m.err != nil {
		// note the empty response
		fprintf(&nr, &err, w, "Error: %v\r\n\r\n\r\n\r\n", m.err)
	} else if m.req != nil {
		fprintf(&nr, &err, w, "\r\n")
		buf, err2 := httputil.DumpRequest(m.req, false)
		if err2 != nil {
			return nr, err2
		}
		write(&nr, &err, w, buf)
	} else if m.resp != nil {
		fprintf(&nr, &err, w, "\r\n")
		buf, err2 := httputil.DumpResponse(m.resp, false)
		if err2 != nil {
			return nr, err2
		}
		write(&nr, &err, w, buf)
	}
	return
}

// HttpLogger is an asynchronous HTTP request/response logger. It traces
// requests and responses headers in a "log" file in logger directory and dumps
// their bodies in files prefixed with the session identifiers.
// Close it to ensure pending items are correctly logged.
type HttpLogger struct {
	path  string
	c     chan *Meta
	errch chan error
}

func NewLogger(basepath string) (*HttpLogger, error) {
	f, err := os.Create(path.Join(basepath, "log"))
	if err != nil {
		return nil, err
	}
	logger := &HttpLogger{basepath, make(chan *Meta), make(chan error)}
	go func() {
		for m := range logger.c {
			if _, err := m.WriteTo(f); err != nil {
				log.Println("Can't write meta", err)
			}
		}
		logger.errch <- f.Close()
	}()
	return logger, nil
}

func (logger *HttpLogger) LogResp(resp *http.Response, ctx *goproxy.ProxyCtx) {
	body := path.Join(logger.path, fmt.Sprintf("%d_resp", ctx.Session))
	from := ""
	if ctx.UserData != nil {
		from = ctx.UserData.(*transport.RoundTripDetails).TCPAddr.String()
	}
	if resp == nil {
		resp = emptyResp
	} else {
		resp.Body = NewTeeReadCloser(resp.Body, NewFileStream(body))
	}
	logger.LogMeta(&Meta{
		resp: resp,
		err:  ctx.Error,
		t:    time.Now(),
		sess: ctx.Session,
		from: from})
}

var emptyResp = &http.Response{}
var emptyReq = &http.Request{}

func (logger *HttpLogger) LogReq(req *http.Request, ctx *goproxy.ProxyCtx) {
	body := path.Join(logger.path, fmt.Sprintf("%d_req", ctx.Session))
	if req == nil {
		req = emptyReq
	} else {
		req.Body = NewTeeReadCloser(req.Body, NewFileStream(body))
	}
	logger.LogMeta(&Meta{
		req:  req,
		err:  ctx.Error,
		t:    time.Now(),
		sess: ctx.Session,
		from: req.RemoteAddr})
}

func (logger *HttpLogger) LogMeta(m *Meta) {
	logger.c <- m
}

func (logger *HttpLogger) Close() error {
	close(logger.c)
	return <-logger.errch
}

// TeeReadCloser extends io.TeeReader by allowing reader and writer to be
// closed.
type TeeReadCloser struct {
	r io.Reader
	w io.WriteCloser
	c io.Closer
}

func NewTeeReadCloser(r io.ReadCloser, w io.WriteCloser) io.ReadCloser {
	return &TeeReadCloser{io.TeeReader(r, w), w, r}
}

func (t *TeeReadCloser) Read(b []byte) (int, error) {
	return t.r.Read(b)
}

// Close attempts to close the reader and write. It returns an error if both
// failed to Close.
func (t *TeeReadCloser) Close() error {
	err1 := t.c.Close()
	err2 := t.w.Close()
	if err1 != nil {
		return err1
	}
	return err2
}

// stoppableListener serves stoppableConn and tracks their lifetime to notify
// when it is safe to terminate the application.
type stoppableListener struct {
	net.Listener
	sync.WaitGroup
}

type stoppableConn struct {
	net.Conn
	wg *sync.WaitGroup
}

func newStoppableListener(l net.Listener) *stoppableListener {
	return &stoppableListener{l, sync.WaitGroup{}}
}

func (sl *stoppableListener) Accept() (net.Conn, error) {
	c, err := sl.Listener.Accept()
	if err != nil {
		return c, err
	}
	sl.Add(1)
	return &stoppableConn{c, &sl.WaitGroup}, nil
}

func (sc *stoppableConn) Close() error {
	sc.wg.Done()
	return sc.Conn.Close()
}

func main() {
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	addr := flag.String("l", ":8080", "on which address should the proxy listen")
	flag.Parse()
	proxy := goproxy.NewProxyHttpServer()
	proxy.Verbose = *verbose
	if err := os.MkdirAll("db", 0755); err != nil {
		log.Fatal("Can't create dir", err)
	}
	logger, err := NewLogger("db")
	if err != nil {
		log.Fatal("can't open log file", err)
	}
	tr := transport.Transport{Proxy: transport.ProxyFromEnvironment}
	// For every incoming request, override the RoundTripper to extract
	// connection information. Store it is session context log it after
	// handling the response.
	proxy.OnRequest().DoFunc(func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		ctx.RoundTripper = goproxy.RoundTripperFunc(func(req *http.Request, ctx *goproxy.ProxyCtx) (resp *http.Response, err error) {
			ctx.UserData, resp, err = tr.DetailedRoundTrip(req)
			return
		})
		logger.LogReq(req, ctx)
		return req, nil
	})
	proxy.OnResponse().DoFunc(func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
		logger.LogResp(resp, ctx)
		return resp
	})
	l, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatal("listen:", err)
	}
	sl := newStoppableListener(l)
	ch := make(chan os.Signal)
	signal.Notify(ch, os.Interrupt)
	go func() {
		<-ch
		log.Println("Got SIGINT exiting")
		sl.Add(1)
		sl.Close()
		logger.Close()
		sl.Done()
	}()
	log.Println("Starting Proxy")
	http.Serve(sl, proxy)
	sl.Wait()
	log.Println("All connections closed - exit")
}
<|endoftext|>


<commit_before>// Copyright 2015 trivago GmbH
//
// 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 format

import (
	"bytes"
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/shared"
)

// StreamMod is a formatter that modifies a message's stream by reading a prefix
// from the message's data (and discarding it).
// The prefix is defined by everything before a given delimiter in the
// message. If no delimiter is found or the prefix is empty the message stream
// is not changed.
// Configuration example
//
//   - "<producer|stream>":
//     Formatter: "format.StreamMod"
//     StreamModFormatter: "format.Forward"
//     StreamModDelimiter: "$"
//
// StreamModFormatter defines the formatter applied after reading the stream.
// This formatter is applied to the data after StreamModDelimiter.
// By default this is set to "format.Forward"
//
// StreamModDelimiter defines the delimiter to search when extracting the stream
// name. By default this is set to ":".
type StreamMod struct {
	base      core.Formatter
	delimiter []byte
}

func init() {
	shared.RuntimeType.Register(StreamMod{})
}

// Configure initializes this formatter with values from a plugin config.
func (format *StreamMod) Configure(conf core.PluginConfig) error {
	plugin, err := core.NewPluginWithType(conf.GetString("StreamModFormatter", "format.Forward"), conf)
	if err != nil {
		return err
	}

	format.delimiter = []byte(conf.GetString("StreamModDelimiter", ":"))
	format.base = plugin.(core.Formatter)

	return nil
}

// Format adds prefix and postfix to the message formatted by the base formatter
func (format *StreamMod) Format(msg core.Message) ([]byte, core.MessageStreamID) {
	modMsg := msg
	prefixEnd := bytes.Index(msg.Data, format.delimiter)

	switch prefixEnd {
	case -1:
	case 0:
		modMsg.Data = msg.Data[1:]
	default:
		firstSpaceIdx := bytes.IndexByte(msg.Data, ' ')
		if (firstSpaceIdx < 0) || (firstSpaceIdx > prefixEnd) {
			modMsg.StreamID = core.GetStreamID(string(msg.Data[:prefixEnd]))
			modMsg.Data = msg.Data[prefixEnd+1:]
		}
	}

	return format.base.Format(modMsg)
}
<commit_msg>Fixed streammod when delimiter is a space Fixed MessageControlProducer flush pointer problems<commit_after>// Copyright 2015 trivago GmbH
//
// 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 format

import (
	"bytes"
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/shared"
)

// StreamMod is a formatter that modifies a message's stream by reading a prefix
// from the message's data (and discarding it).
// The prefix is defined by everything before a given delimiter in the
// message. If no delimiter is found or the prefix is empty the message stream
// is not changed.
// Configuration example
//
//   - "<producer|stream>":
//     Formatter: "format.StreamMod"
//     StreamModFormatter: "format.Forward"
//     StreamModDelimiter: "$"
//
// StreamModFormatter defines the formatter applied after reading the stream.
// This formatter is applied to the data after StreamModDelimiter.
// By default this is set to "format.Forward"
//
// StreamModDelimiter defines the delimiter to search when extracting the stream
// name. By default this is set to ":".
type StreamMod struct {
	base      core.Formatter
	delimiter []byte
}

func init() {
	shared.RuntimeType.Register(StreamMod{})
}

// Configure initializes this formatter with values from a plugin config.
func (format *StreamMod) Configure(conf core.PluginConfig) error {
	plugin, err := core.NewPluginWithType(conf.GetString("StreamModFormatter", "format.Forward"), conf)
	if err != nil {
		return err
	}

	format.delimiter = []byte(conf.GetString("StreamModDelimiter", ":"))
	format.base = plugin.(core.Formatter)

	return nil
}

// Format adds prefix and postfix to the message formatted by the base formatter
func (format *StreamMod) Format(msg core.Message) ([]byte, core.MessageStreamID) {
	modMsg := msg
	prefixEnd := bytes.Index(msg.Data, format.delimiter)

	switch prefixEnd {
	case -1:
	case 0:
		modMsg.Data = msg.Data[1:]
	default:
		firstSpaceIdx := bytes.IndexByte(msg.Data, ' ')
		if (firstSpaceIdx < 0) || (firstSpaceIdx >= prefixEnd) {
			modMsg.StreamID = core.GetStreamID(string(msg.Data[:prefixEnd]))
			modMsg.Data = msg.Data[prefixEnd+1:]
		}
	}

	return format.base.Format(modMsg)
}
<|endoftext|>


<commit_before>package handler

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

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"

	"github.com/Lunchr/luncher-api/db"
	"github.com/Lunchr/luncher-api/db/model"
	"github.com/Lunchr/luncher-api/router"
	"github.com/Lunchr/luncher-api/session"
	"github.com/deiwin/facebook"
	"github.com/julienschmidt/httprouter"
)

// OfferGroupPost handles GET requests to /restaurant/posts/:date. It returns all current day's offers for the region.
func OfferGroupPost(c db.OfferGroupPosts, sessionManager session.Manager, users db.Users, restaurants db.Restaurants) router.HandlerWithParams {
	handler := func(w http.ResponseWriter, r *http.Request, user *model.User, restaurant *model.Restaurant,
		date model.DateWithoutTime) *router.HandlerError {
		post, err := c.GetByDate(date, restaurant.ID)
		if err == mgo.ErrNotFound {
			return router.NewHandlerError(err, "Offer group post not found", http.StatusNotFound)
		} else if err != nil {
			return router.NewHandlerError(err, "An error occured while trying to fetch a offer group post", http.StatusInternalServerError)
		}
		return writeJSON(w, post)
	}
	return forDate(sessionManager, users, restaurants, handler)
}

// PostOfferGroupPost handles POST requests to /restaurant/posts. It stores the info in the DB and updates the post in FB.
func PostOfferGroupPost(c db.OfferGroupPosts, sessionManager session.Manager, users db.Users, restaurants db.Restaurants,
	offers db.Offers, regions db.Regions, fbAuth facebook.Authenticator) router.Handler {
	handler := func(w http.ResponseWriter, r *http.Request, user *model.User, restaurant *model.Restaurant) *router.HandlerError {
		post, handlerErr := parseOfferGroupPost(r, restaurant)
		if handlerErr != nil {
			return handlerErr
		}
		insertedPosts, err := c.Insert(post)
		if err != nil {
			return router.NewHandlerError(err, "Failed to store the post in the DB", http.StatusInternalServerError)
		}
		insertedPost := insertedPosts[0]
		if handlerErr = updateGroupPost(insertedPost, user, restaurant, offers, regions, c, fbAuth); handlerErr != nil {
			return handlerErr
		}
		return writeJSON(w, insertedPost)
	}
	return forRestaurant(sessionManager, users, restaurants, handler)
}

// PutOfferGroupPost handles PUT requests to /restaurant/posts/:date. It stores the info in the DB and updates the post in FB.
func PutOfferGroupPost(c db.OfferGroupPosts, sessionManager session.Manager, users db.Users, restaurants db.Restaurants,
	offers db.Offers, regions db.Regions, fbAuth facebook.Authenticator) router.HandlerWithParams {
	handler := func(w http.ResponseWriter, r *http.Request, user *model.User, restaurant *model.Restaurant,
		date model.DateWithoutTime) *router.HandlerError {
		post, handlerErr := parseOfferGroupPost(r, restaurant)
		if handlerErr != nil {
			return handlerErr
		}
		if post.Date != date {
			return router.NewSimpleHandlerError("Unexpected date value", http.StatusBadRequest)
		}
		err := c.UpdateByID(post.ID, post)
		if err != nil {
			return router.NewSimpleHandlerError("Failed to insert the post to DB", http.StatusBadRequest)
		}
		return writeJSON(w, post)
	}
	return forDate(sessionManager, users, restaurants, handler)
}

type HandlerWithRestaurantAndDate func(w http.ResponseWriter, r *http.Request, user *model.User, restaurant *model.Restaurant,
	date model.DateWithoutTime) *router.HandlerError

func forDate(sessionManager session.Manager, users db.Users, restaurants db.Restaurants,
	handler HandlerWithRestaurantAndDate) router.HandlerWithParams {
	handlerWithRestaurant := func(w http.ResponseWriter, r *http.Request, ps httprouter.Params, user *model.User,
		restaurant *model.Restaurant) *router.HandlerError {
		date := model.DateWithoutTime(ps.ByName("date"))
		if date == "" {
			return router.NewStringHandlerError("Date not specified!", "Please specify a date", http.StatusBadRequest)
		}
		if !date.IsValid() {
			return router.NewSimpleHandlerError("Invalid date specified", http.StatusBadRequest)
		}
		return handler(w, r, user, restaurant, date)
	}
	return forRestaurantWithParams(sessionManager, users, restaurants, handlerWithRestaurant)
}

func parseOfferGroupPost(r *http.Request, restaurant *model.Restaurant) (*model.OfferGroupPost, *router.HandlerError) {
	var post struct {
		ID              bson.ObjectId `json:"_id"`
		MessageTemplate string        `json:"message_template"`
		Date            string        `json:"date"`
	}
	err := json.NewDecoder(r.Body).Decode(&post)
	if err != nil {
		return nil, router.NewHandlerError(err, "Failed to parse the post", http.StatusBadRequest)
	}
	date := model.DateWithoutTime(post.Date)
	if date == "" {
		return nil, router.NewStringHandlerError("Date not specified!", "Please specify a date", http.StatusBadRequest)
	}
	if !date.IsValid() {
		return nil, router.NewSimpleHandlerError("Invalid date specified", http.StatusBadRequest)
	}
	return &model.OfferGroupPost{
		ID:              post.ID,
		MessageTemplate: post.MessageTemplate,
		Date:            date,
		RestaurantID:    restaurant.ID,
	}, nil
}

func updateGroupPostForDate(date model.DateWithoutTime, user *model.User, restaurant *model.Restaurant, offers db.Offers,
	regions db.Regions, groupPosts db.OfferGroupPosts, fbAuth facebook.Authenticator) *router.HandlerError {
	post, err := groupPosts.GetByDate(date, restaurant.ID)
	if err == mgo.ErrNotFound {
		postToInsert := &model.OfferGroupPost{
			RestaurantID:    restaurant.ID,
			Date:            date,
			MessageTemplate: restaurant.DefaultGroupPostMessageTemplate,
		}
		insertedPosts, err := groupPosts.Insert(postToInsert)
		if err != nil {
			router.NewHandlerError(err, "Failed to create a group post with restaurant defaults", http.StatusInternalServerError)
		}
		post = insertedPosts[0]
	} else if err != nil {
		return router.NewHandlerError(err, "Failed to fetch a group post for that date", http.StatusInternalServerError)
	}
	return updateGroupPost(post, user, restaurant, offers, regions, groupPosts, fbAuth)
}

func updateGroupPost(post *model.OfferGroupPost, user *model.User, restaurant *model.Restaurant, offers db.Offers,
	regions db.Regions, groupPosts db.OfferGroupPosts, fbAuth facebook.Authenticator) *router.HandlerError {
	if restaurant.FacebookPageID == "" {
		return nil
	}
	fbAPI := fbAuth.APIConnection(&user.Session.FacebookUserToken)
	// Remove the current post from FB, if it's already there
	if post.FBPostID != "" {
		err := fbAPI.PostDelete(user.Session.FacebookPageToken, post.FBPostID)
		if err != nil {
			return router.NewHandlerError(err, "Failed to delete the current post from Facebook", http.StatusBadGateway)
		}
	}
	offersForDate, handlerErr := getOffersForDate(post.Date, restaurant, offers, regions)
	if handlerErr != nil {
		return handlerErr
	} else if len(offersForDate) == 0 {
		return nil
	}
	message := formFBMessage(post, offersForDate)
	// Add the new version
	fbPost, err := fbAPI.PagePublish(user.Session.FacebookPageToken, restaurant.FacebookPageID, message)
	if err != nil {
		return router.NewHandlerError(err, "Failed to post the offer to Facebook", http.StatusBadGateway)
	}
	post.FBPostID = fbPost.ID
	if err = groupPosts.UpdateByID(post.ID, post); err != nil {
		return router.NewHandlerError(err, "Failed to update a group post in the DB", http.StatusInternalServerError)
	}
	return nil
}

func formFBMessage(post *model.OfferGroupPost, offers []*model.Offer) string {
	offerMessages := make([]string, len(offers))
	for i, offer := range offers {
		offerMessages[i] = formFBOfferMessage(offer)
	}
	offersMessage := strings.Join(offerMessages, "\n")
	return fmt.Sprintf("%s\n\n%s", post.MessageTemplate, offersMessage)
}

func formFBOfferMessage(o *model.Offer) string {
	return fmt.Sprintf("%s - %.2f€", o.Title, o.Price)
}

func getOffersForDate(date model.DateWithoutTime, restaurant *model.Restaurant, offers db.Offers, regions db.Regions) ([]*model.Offer, *router.HandlerError) {
	region, err := regions.GetName(restaurant.Region)
	if err != nil {
		return nil, router.NewHandlerError(err, "Failed to find the restaurant's region", http.StatusInternalServerError)
	}
	location, err := time.LoadLocation(region.Location)
	if err != nil {
		return nil, router.NewHandlerError(err, "Failed to load region's location", http.StatusInternalServerError)
	}
	startTime, endTime, err := date.TimeBounds(location)
	if err != nil {
		return nil, router.NewHandlerError(err, "Failed to parse a date", http.StatusInternalServerError)
	}
	offersForDate, err := offers.GetForRestaurantWithinTimeBounds(restaurant.ID, startTime, endTime)
	if err != nil {
		return nil, router.NewHandlerError(err, "Failed to find offers for this date", http.StatusInternalServerError)
	}
	return offersForDate, nil
}
<commit_msg>Add a TODO for a hard-coded currency symbol<commit_after>package handler

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

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"

	"github.com/Lunchr/luncher-api/db"
	"github.com/Lunchr/luncher-api/db/model"
	"github.com/Lunchr/luncher-api/router"
	"github.com/Lunchr/luncher-api/session"
	"github.com/deiwin/facebook"
	"github.com/julienschmidt/httprouter"
)

// OfferGroupPost handles GET requests to /restaurant/posts/:date. It returns all current day's offers for the region.
func OfferGroupPost(c db.OfferGroupPosts, sessionManager session.Manager, users db.Users, restaurants db.Restaurants) router.HandlerWithParams {
	handler := func(w http.ResponseWriter, r *http.Request, user *model.User, restaurant *model.Restaurant,
		date model.DateWithoutTime) *router.HandlerError {
		post, err := c.GetByDate(date, restaurant.ID)
		if err == mgo.ErrNotFound {
			return router.NewHandlerError(err, "Offer group post not found", http.StatusNotFound)
		} else if err != nil {
			return router.NewHandlerError(err, "An error occured while trying to fetch a offer group post", http.StatusInternalServerError)
		}
		return writeJSON(w, post)
	}
	return forDate(sessionManager, users, restaurants, handler)
}

// PostOfferGroupPost handles POST requests to /restaurant/posts. It stores the info in the DB and updates the post in FB.
func PostOfferGroupPost(c db.OfferGroupPosts, sessionManager session.Manager, users db.Users, restaurants db.Restaurants,
	offers db.Offers, regions db.Regions, fbAuth facebook.Authenticator) router.Handler {
	handler := func(w http.ResponseWriter, r *http.Request, user *model.User, restaurant *model.Restaurant) *router.HandlerError {
		post, handlerErr := parseOfferGroupPost(r, restaurant)
		if handlerErr != nil {
			return handlerErr
		}
		insertedPosts, err := c.Insert(post)
		if err != nil {
			return router.NewHandlerError(err, "Failed to store the post in the DB", http.StatusInternalServerError)
		}
		insertedPost := insertedPosts[0]
		if handlerErr = updateGroupPost(insertedPost, user, restaurant, offers, regions, c, fbAuth); handlerErr != nil {
			return handlerErr
		}
		return writeJSON(w, insertedPost)
	}
	return forRestaurant(sessionManager, users, restaurants, handler)
}

// PutOfferGroupPost handles PUT requests to /restaurant/posts/:date. It stores the info in the DB and updates the post in FB.
func PutOfferGroupPost(c db.OfferGroupPosts, sessionManager session.Manager, users db.Users, restaurants db.Restaurants,
	offers db.Offers, regions db.Regions, fbAuth facebook.Authenticator) router.HandlerWithParams {
	handler := func(w http.ResponseWriter, r *http.Request, user *model.User, restaurant *model.Restaurant,
		date model.DateWithoutTime) *router.HandlerError {
		post, handlerErr := parseOfferGroupPost(r, restaurant)
		if handlerErr != nil {
			return handlerErr
		}
		if post.Date != date {
			return router.NewSimpleHandlerError("Unexpected date value", http.StatusBadRequest)
		}
		err := c.UpdateByID(post.ID, post)
		if err != nil {
			return router.NewSimpleHandlerError("Failed to insert the post to DB", http.StatusBadRequest)
		}
		return writeJSON(w, post)
	}
	return forDate(sessionManager, users, restaurants, handler)
}

type HandlerWithRestaurantAndDate func(w http.ResponseWriter, r *http.Request, user *model.User, restaurant *model.Restaurant,
	date model.DateWithoutTime) *router.HandlerError

func forDate(sessionManager session.Manager, users db.Users, restaurants db.Restaurants,
	handler HandlerWithRestaurantAndDate) router.HandlerWithParams {
	handlerWithRestaurant := func(w http.ResponseWriter, r *http.Request, ps httprouter.Params, user *model.User,
		restaurant *model.Restaurant) *router.HandlerError {
		date := model.DateWithoutTime(ps.ByName("date"))
		if date == "" {
			return router.NewStringHandlerError("Date not specified!", "Please specify a date", http.StatusBadRequest)
		}
		if !date.IsValid() {
			return router.NewSimpleHandlerError("Invalid date specified", http.StatusBadRequest)
		}
		return handler(w, r, user, restaurant, date)
	}
	return forRestaurantWithParams(sessionManager, users, restaurants, handlerWithRestaurant)
}

func parseOfferGroupPost(r *http.Request, restaurant *model.Restaurant) (*model.OfferGroupPost, *router.HandlerError) {
	var post struct {
		ID              bson.ObjectId `json:"_id"`
		MessageTemplate string        `json:"message_template"`
		Date            string        `json:"date"`
	}
	err := json.NewDecoder(r.Body).Decode(&post)
	if err != nil {
		return nil, router.NewHandlerError(err, "Failed to parse the post", http.StatusBadRequest)
	}
	date := model.DateWithoutTime(post.Date)
	if date == "" {
		return nil, router.NewStringHandlerError("Date not specified!", "Please specify a date", http.StatusBadRequest)
	}
	if !date.IsValid() {
		return nil, router.NewSimpleHandlerError("Invalid date specified", http.StatusBadRequest)
	}
	return &model.OfferGroupPost{
		ID:              post.ID,
		MessageTemplate: post.MessageTemplate,
		Date:            date,
		RestaurantID:    restaurant.ID,
	}, nil
}

func updateGroupPostForDate(date model.DateWithoutTime, user *model.User, restaurant *model.Restaurant, offers db.Offers,
	regions db.Regions, groupPosts db.OfferGroupPosts, fbAuth facebook.Authenticator) *router.HandlerError {
	post, err := groupPosts.GetByDate(date, restaurant.ID)
	if err == mgo.ErrNotFound {
		postToInsert := &model.OfferGroupPost{
			RestaurantID:    restaurant.ID,
			Date:            date,
			MessageTemplate: restaurant.DefaultGroupPostMessageTemplate,
		}
		insertedPosts, err := groupPosts.Insert(postToInsert)
		if err != nil {
			router.NewHandlerError(err, "Failed to create a group post with restaurant defaults", http.StatusInternalServerError)
		}
		post = insertedPosts[0]
	} else if err != nil {
		return router.NewHandlerError(err, "Failed to fetch a group post for that date", http.StatusInternalServerError)
	}
	return updateGroupPost(post, user, restaurant, offers, regions, groupPosts, fbAuth)
}

func updateGroupPost(post *model.OfferGroupPost, user *model.User, restaurant *model.Restaurant, offers db.Offers,
	regions db.Regions, groupPosts db.OfferGroupPosts, fbAuth facebook.Authenticator) *router.HandlerError {
	if restaurant.FacebookPageID == "" {
		return nil
	}
	fbAPI := fbAuth.APIConnection(&user.Session.FacebookUserToken)
	// Remove the current post from FB, if it's already there
	if post.FBPostID != "" {
		err := fbAPI.PostDelete(user.Session.FacebookPageToken, post.FBPostID)
		if err != nil {
			return router.NewHandlerError(err, "Failed to delete the current post from Facebook", http.StatusBadGateway)
		}
	}
	offersForDate, handlerErr := getOffersForDate(post.Date, restaurant, offers, regions)
	if handlerErr != nil {
		return handlerErr
	} else if len(offersForDate) == 0 {
		return nil
	}
	message := formFBMessage(post, offersForDate)
	// Add the new version
	fbPost, err := fbAPI.PagePublish(user.Session.FacebookPageToken, restaurant.FacebookPageID, message)
	if err != nil {
		return router.NewHandlerError(err, "Failed to post the offer to Facebook", http.StatusBadGateway)
	}
	post.FBPostID = fbPost.ID
	if err = groupPosts.UpdateByID(post.ID, post); err != nil {
		return router.NewHandlerError(err, "Failed to update a group post in the DB", http.StatusInternalServerError)
	}
	return nil
}

func formFBMessage(post *model.OfferGroupPost, offers []*model.Offer) string {
	offerMessages := make([]string, len(offers))
	for i, offer := range offers {
		offerMessages[i] = formFBOfferMessage(offer)
	}
	offersMessage := strings.Join(offerMessages, "\n")
	return fmt.Sprintf("%s\n\n%s", post.MessageTemplate, offersMessage)
}

func formFBOfferMessage(o *model.Offer) string {
	// TODO get rid of the hard-coded €
	return fmt.Sprintf("%s - %.2f€", o.Title, o.Price)
}

func getOffersForDate(date model.DateWithoutTime, restaurant *model.Restaurant, offers db.Offers, regions db.Regions) ([]*model.Offer, *router.HandlerError) {
	region, err := regions.GetName(restaurant.Region)
	if err != nil {
		return nil, router.NewHandlerError(err, "Failed to find the restaurant's region", http.StatusInternalServerError)
	}
	location, err := time.LoadLocation(region.Location)
	if err != nil {
		return nil, router.NewHandlerError(err, "Failed to load region's location", http.StatusInternalServerError)
	}
	startTime, endTime, err := date.TimeBounds(location)
	if err != nil {
		return nil, router.NewHandlerError(err, "Failed to parse a date", http.StatusInternalServerError)
	}
	offersForDate, err := offers.GetForRestaurantWithinTimeBounds(restaurant.ID, startTime, endTime)
	if err != nil {
		return nil, router.NewHandlerError(err, "Failed to find offers for this date", http.StatusInternalServerError)
	}
	return offersForDate, nil
}
<|endoftext|>


<commit_before>// Package fluentd provides the log driver for forwarding server logs
// to fluentd endpoints.
package fluentd // import "github.com/docker/docker/daemon/logger/fluentd"

import (
	"fmt"
	"math"
	"net"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/docker/docker/daemon/logger"
	"github.com/docker/docker/daemon/logger/loggerutils"
	"github.com/docker/docker/errdefs"
	"github.com/docker/docker/pkg/urlutil"
	units "github.com/docker/go-units"
	"github.com/fluent/fluent-logger-golang/fluent"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

type fluentd struct {
	tag           string
	containerID   string
	containerName string
	writer        *fluent.Fluent
	extra         map[string]string
}

type location struct {
	protocol string
	host     string
	port     int
	path     string
}

const (
	name = "fluentd"

	defaultProtocol    = "tcp"
	defaultHost        = "127.0.0.1"
	defaultPort        = 24224
	defaultBufferLimit = 1024 * 1024

	// logger tries to reconnect 2**32 - 1 times
	// failed (and panic) after 204 years [ 1.5 ** (2**32 - 1) - 1 seconds]
	defaultRetryWait  = 1000
	defaultMaxRetries = math.MaxInt32

	addressKey            = "fluentd-address"
	bufferLimitKey        = "fluentd-buffer-limit"
	retryWaitKey          = "fluentd-retry-wait"
	maxRetriesKey         = "fluentd-max-retries"
	asyncConnectKey       = "fluentd-async-connect"
	subSecondPrecisionKey = "fluentd-sub-second-precision"
)

func init() {
	if err := logger.RegisterLogDriver(name, New); err != nil {
		logrus.Fatal(err)
	}
	if err := logger.RegisterLogOptValidator(name, ValidateLogOpt); err != nil {
		logrus.Fatal(err)
	}
}

// New creates a fluentd logger using the configuration passed in on
// the context. The supported context configuration variable is
// fluentd-address.
func New(info logger.Info) (logger.Logger, error) {
	loc, err := parseAddress(info.Config[addressKey])
	if err != nil {
		return nil, errdefs.InvalidParameter(err)
	}

	tag, err := loggerutils.ParseLogTag(info, loggerutils.DefaultTemplate)
	if err != nil {
		return nil, errdefs.InvalidParameter(err)
	}

	extra, err := info.ExtraAttributes(nil)
	if err != nil {
		return nil, errdefs.InvalidParameter(err)
	}

	bufferLimit := defaultBufferLimit
	if info.Config[bufferLimitKey] != "" {
		bl64, err := units.RAMInBytes(info.Config[bufferLimitKey])
		if err != nil {
			return nil, err
		}
		bufferLimit = int(bl64)
	}

	retryWait := defaultRetryWait
	if info.Config[retryWaitKey] != "" {
		rwd, err := time.ParseDuration(info.Config[retryWaitKey])
		if err != nil {
			return nil, err
		}
		retryWait = int(rwd.Seconds() * 1000)
	}

	maxRetries := defaultMaxRetries
	if info.Config[maxRetriesKey] != "" {
		mr64, err := strconv.ParseUint(info.Config[maxRetriesKey], 10, strconv.IntSize)
		if err != nil {
			return nil, err
		}
		maxRetries = int(mr64)
	}

	asyncConnect := false
	if info.Config[asyncConnectKey] != "" {
		if asyncConnect, err = strconv.ParseBool(info.Config[asyncConnectKey]); err != nil {
			return nil, err
		}
	}

	subSecondPrecision := false
	if info.Config[subSecondPrecisionKey] != "" {
		if subSecondPrecision, err = strconv.ParseBool(info.Config[subSecondPrecisionKey]); err != nil {
			return nil, err
		}
	}

	fluentConfig := fluent.Config{
		FluentPort:         loc.port,
		FluentHost:         loc.host,
		FluentNetwork:      loc.protocol,
		FluentSocketPath:   loc.path,
		BufferLimit:        bufferLimit,
		RetryWait:          retryWait,
		MaxRetry:           maxRetries,
		Async:              asyncConnect,
		SubSecondPrecision: subSecondPrecision,
	}

	logrus.WithField("container", info.ContainerID).WithField("config", fluentConfig).
		Debug("logging driver fluentd configured")

	log, err := fluent.New(fluentConfig)
	if err != nil {
		return nil, err
	}
	return &fluentd{
		tag:           tag,
		containerID:   info.ContainerID,
		containerName: info.ContainerName,
		writer:        log,
		extra:         extra,
	}, nil
}

func (f *fluentd) Log(msg *logger.Message) error {
	data := map[string]string{
		"container_id":   f.containerID,
		"container_name": f.containerName,
		"source":         msg.Source,
		"log":            string(msg.Line),
	}
	for k, v := range f.extra {
		data[k] = v
	}
	if msg.PLogMetaData != nil {
		data["partial_message"] = "true"
		data["partial_id"] = msg.PLogMetaData.ID
		data["partial_ordinal"] = strconv.Itoa(msg.PLogMetaData.Ordinal)
		data["partial_last"] = strconv.FormatBool(msg.PLogMetaData.Last)
	}

	ts := msg.Timestamp
	logger.PutMessage(msg)
	// fluent-logger-golang buffers logs from failures and disconnections,
	// and these are transferred again automatically.
	return f.writer.PostWithTime(f.tag, ts, data)
}

func (f *fluentd) Close() error {
	return f.writer.Close()
}

func (f *fluentd) Name() string {
	return name
}

// ValidateLogOpt looks for fluentd specific log option fluentd-address.
func ValidateLogOpt(cfg map[string]string) error {
	for key := range cfg {
		switch key {
		case "env":
		case "env-regex":
		case "labels":
		case "labels-regex":
		case "tag":
		case addressKey:
		case bufferLimitKey:
		case retryWaitKey:
		case maxRetriesKey:
		case asyncConnectKey:
		case subSecondPrecisionKey:
			// Accepted
		default:
			return fmt.Errorf("unknown log opt '%s' for fluentd log driver", key)
		}
	}

	_, err := parseAddress(cfg[addressKey])
	return err
}

func parseAddress(address string) (*location, error) {
	if address == "" {
		return &location{
			protocol: defaultProtocol,
			host:     defaultHost,
			port:     defaultPort,
			path:     "",
		}, nil
	}

	protocol := defaultProtocol
	givenAddress := address
	if urlutil.IsTransportURL(address) {
		url, err := url.Parse(address)
		if err != nil {
			return nil, errors.Wrapf(err, "invalid fluentd-address %s", givenAddress)
		}
		// unix and unixgram socket
		if url.Scheme == "unix" || url.Scheme == "unixgram" {
			return &location{
				protocol: url.Scheme,
				host:     "",
				port:     0,
				path:     url.Path,
			}, nil
		}
		// tcp|udp
		protocol = url.Scheme
		address = url.Host
	}

	host, port, err := net.SplitHostPort(address)
	if err != nil {
		if !strings.Contains(err.Error(), "missing port in address") {
			return nil, errors.Wrapf(err, "invalid fluentd-address %s", givenAddress)
		}
		return &location{
			protocol: protocol,
			host:     host,
			port:     defaultPort,
			path:     "",
		}, nil
	}

	portnum, err := strconv.Atoi(port)
	if err != nil {
		return nil, errors.Wrapf(err, "invalid fluentd-address %s", givenAddress)
	}
	return &location{
		protocol: protocol,
		host:     host,
		port:     portnum,
		path:     "",
	}, nil
}
<commit_msg>Fluentd: sort consts alphabetically<commit_after>// Package fluentd provides the log driver for forwarding server logs
// to fluentd endpoints.
package fluentd // import "github.com/docker/docker/daemon/logger/fluentd"

import (
	"fmt"
	"math"
	"net"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/docker/docker/daemon/logger"
	"github.com/docker/docker/daemon/logger/loggerutils"
	"github.com/docker/docker/errdefs"
	"github.com/docker/docker/pkg/urlutil"
	units "github.com/docker/go-units"
	"github.com/fluent/fluent-logger-golang/fluent"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

type fluentd struct {
	tag           string
	containerID   string
	containerName string
	writer        *fluent.Fluent
	extra         map[string]string
}

type location struct {
	protocol string
	host     string
	port     int
	path     string
}

const (
	name = "fluentd"

	defaultBufferLimit = 1024 * 1024
	defaultHost        = "127.0.0.1"
	defaultPort        = 24224
	defaultProtocol    = "tcp"

	// logger tries to reconnect 2**32 - 1 times
	// failed (and panic) after 204 years [ 1.5 ** (2**32 - 1) - 1 seconds]
	defaultMaxRetries = math.MaxInt32
	defaultRetryWait  = 1000

	addressKey            = "fluentd-address"
	asyncConnectKey       = "fluentd-async-connect"
	bufferLimitKey        = "fluentd-buffer-limit"
	maxRetriesKey         = "fluentd-max-retries"
	retryWaitKey          = "fluentd-retry-wait"
	subSecondPrecisionKey = "fluentd-sub-second-precision"
)

func init() {
	if err := logger.RegisterLogDriver(name, New); err != nil {
		logrus.Fatal(err)
	}
	if err := logger.RegisterLogOptValidator(name, ValidateLogOpt); err != nil {
		logrus.Fatal(err)
	}
}

// New creates a fluentd logger using the configuration passed in on
// the context. The supported context configuration variable is
// fluentd-address.
func New(info logger.Info) (logger.Logger, error) {
	loc, err := parseAddress(info.Config[addressKey])
	if err != nil {
		return nil, errdefs.InvalidParameter(err)
	}

	tag, err := loggerutils.ParseLogTag(info, loggerutils.DefaultTemplate)
	if err != nil {
		return nil, errdefs.InvalidParameter(err)
	}

	extra, err := info.ExtraAttributes(nil)
	if err != nil {
		return nil, errdefs.InvalidParameter(err)
	}

	bufferLimit := defaultBufferLimit
	if info.Config[bufferLimitKey] != "" {
		bl64, err := units.RAMInBytes(info.Config[bufferLimitKey])
		if err != nil {
			return nil, err
		}
		bufferLimit = int(bl64)
	}

	retryWait := defaultRetryWait
	if info.Config[retryWaitKey] != "" {
		rwd, err := time.ParseDuration(info.Config[retryWaitKey])
		if err != nil {
			return nil, err
		}
		retryWait = int(rwd.Seconds() * 1000)
	}

	maxRetries := defaultMaxRetries
	if info.Config[maxRetriesKey] != "" {
		mr64, err := strconv.ParseUint(info.Config[maxRetriesKey], 10, strconv.IntSize)
		if err != nil {
			return nil, err
		}
		maxRetries = int(mr64)
	}

	asyncConnect := false
	if info.Config[asyncConnectKey] != "" {
		if asyncConnect, err = strconv.ParseBool(info.Config[asyncConnectKey]); err != nil {
			return nil, err
		}
	}

	subSecondPrecision := false
	if info.Config[subSecondPrecisionKey] != "" {
		if subSecondPrecision, err = strconv.ParseBool(info.Config[subSecondPrecisionKey]); err != nil {
			return nil, err
		}
	}

	fluentConfig := fluent.Config{
		FluentPort:         loc.port,
		FluentHost:         loc.host,
		FluentNetwork:      loc.protocol,
		FluentSocketPath:   loc.path,
		BufferLimit:        bufferLimit,
		RetryWait:          retryWait,
		MaxRetry:           maxRetries,
		Async:              asyncConnect,
		SubSecondPrecision: subSecondPrecision,
	}

	logrus.WithField("container", info.ContainerID).WithField("config", fluentConfig).
		Debug("logging driver fluentd configured")

	log, err := fluent.New(fluentConfig)
	if err != nil {
		return nil, err
	}
	return &fluentd{
		tag:           tag,
		containerID:   info.ContainerID,
		containerName: info.ContainerName,
		writer:        log,
		extra:         extra,
	}, nil
}

func (f *fluentd) Log(msg *logger.Message) error {
	data := map[string]string{
		"container_id":   f.containerID,
		"container_name": f.containerName,
		"source":         msg.Source,
		"log":            string(msg.Line),
	}
	for k, v := range f.extra {
		data[k] = v
	}
	if msg.PLogMetaData != nil {
		data["partial_message"] = "true"
		data["partial_id"] = msg.PLogMetaData.ID
		data["partial_ordinal"] = strconv.Itoa(msg.PLogMetaData.Ordinal)
		data["partial_last"] = strconv.FormatBool(msg.PLogMetaData.Last)
	}

	ts := msg.Timestamp
	logger.PutMessage(msg)
	// fluent-logger-golang buffers logs from failures and disconnections,
	// and these are transferred again automatically.
	return f.writer.PostWithTime(f.tag, ts, data)
}

func (f *fluentd) Close() error {
	return f.writer.Close()
}

func (f *fluentd) Name() string {
	return name
}

// ValidateLogOpt looks for fluentd specific log option fluentd-address.
func ValidateLogOpt(cfg map[string]string) error {
	for key := range cfg {
		switch key {
		case "env":
		case "env-regex":
		case "labels":
		case "labels-regex":
		case "tag":

		case addressKey:
		case asyncConnectKey:
		case bufferLimitKey:
		case maxRetriesKey:
		case retryWaitKey:
		case subSecondPrecisionKey:
			// Accepted
		default:
			return fmt.Errorf("unknown log opt '%s' for fluentd log driver", key)
		}
	}

	_, err := parseAddress(cfg[addressKey])
	return err
}

func parseAddress(address string) (*location, error) {
	if address == "" {
		return &location{
			protocol: defaultProtocol,
			host:     defaultHost,
			port:     defaultPort,
			path:     "",
		}, nil
	}

	protocol := defaultProtocol
	givenAddress := address
	if urlutil.IsTransportURL(address) {
		url, err := url.Parse(address)
		if err != nil {
			return nil, errors.Wrapf(err, "invalid fluentd-address %s", givenAddress)
		}
		// unix and unixgram socket
		if url.Scheme == "unix" || url.Scheme == "unixgram" {
			return &location{
				protocol: url.Scheme,
				host:     "",
				port:     0,
				path:     url.Path,
			}, nil
		}
		// tcp|udp
		protocol = url.Scheme
		address = url.Host
	}

	host, port, err := net.SplitHostPort(address)
	if err != nil {
		if !strings.Contains(err.Error(), "missing port in address") {
			return nil, errors.Wrapf(err, "invalid fluentd-address %s", givenAddress)
		}
		return &location{
			protocol: protocol,
			host:     host,
			port:     defaultPort,
			path:     "",
		}, nil
	}

	portnum, err := strconv.Atoi(port)
	if err != nil {
		return nil, errors.Wrapf(err, "invalid fluentd-address %s", givenAddress)
	}
	return &location{
		protocol: protocol,
		host:     host,
		port:     portnum,
		path:     "",
	}, nil
}
<|endoftext|>


<commit_before>// Copyright 2022 Google LLC
//
// 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 e2e

import (
	"fmt"
	"testing"
	"time"

	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	"kpt.dev/configsync/e2e/nomostest"
	"kpt.dev/configsync/e2e/nomostest/metrics"
	"kpt.dev/configsync/e2e/nomostest/ntopts"
	nomostesting "kpt.dev/configsync/e2e/nomostest/testing"
	"kpt.dev/configsync/pkg/api/configsync"
	"kpt.dev/configsync/pkg/core"
	"kpt.dev/configsync/pkg/syncer/differ"
	"kpt.dev/configsync/pkg/testing/fake"
	"kpt.dev/configsync/pkg/util"
	"sigs.k8s.io/cli-utils/pkg/common"
)

var preventDeletion = core.Annotation(common.LifecycleDeleteAnnotation, common.PreventDeletion)

func TestPreventDeletionNamespace(t *testing.T) {
	nt := nomostest.New(t, nomostesting.Lifecycle)

	// Ensure the Namespace doesn't already exist.
	err := nt.ValidateNotFound("shipping", "", &corev1.Namespace{})
	if err != nil {
		nt.T.Fatal(err)
	}

	role := fake.RoleObject(core.Name("shipping-admin"))
	role.Rules = []rbacv1.PolicyRule{{
		APIGroups: []string{""},
		Resources: []string{"configmaps"},
		Verbs:     []string{"get"},
	}}

	// Declare the Namespace with the lifecycle annotation, and ensure it is created.
	nt.RootRepos[configsync.RootSyncName].Add("acme/namespaces/shipping/ns.yaml",
		fake.NamespaceObject("shipping", preventDeletion))
	nt.RootRepos[configsync.RootSyncName].Add("acme/namespaces/shipping/role.yaml", role)
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("declare Namespace with prevent deletion lifecycle annotation")
	nt.WaitForRepoSyncs()

	err = nt.Validate("shipping", "", &corev1.Namespace{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Validate multi-repo metrics.
	err = nt.ValidateMetrics(nomostest.SyncMetricsToLatestCommit(nt), func() error {
		err := nt.ValidateMultiRepoMetrics(nomostest.DefaultRootReconcilerName,
			nt.DefaultRootSyncObjectCount()+2, // 2 for the test Namespace & Role
			metrics.ResourceCreated("Namespace"), metrics.ResourceCreated("Role"))
		if err != nil {
			return err
		}
		return nt.ValidateErrorMetricsNotFound()
	})
	if err != nil {
		nt.T.Error(err)
	}

	// Delete the declaration and ensure the Namespace isn't deleted.
	nt.RootRepos[configsync.RootSyncName].Remove("acme/namespaces/shipping/ns.yaml")
	nt.RootRepos[configsync.RootSyncName].Remove("acme/namespaces/shipping/role.yaml")
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove Namespace shipping declaration")
	nt.WaitForRepoSyncs()

	// Ensure we kept the undeclared Namespace that had the "deletion: prevent" annotation.
	err = nt.Validate("shipping", "", &corev1.Namespace{},
		nomostest.NotPendingDeletion)
	if err != nil {
		nt.T.Fatal(err)
	}
	// Ensure we deleted the undeclared Role that doesn't have the annotation.
	err = nt.ValidateNotFound("shipping-admin", "shipping", &rbacv1.Role{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Validate multi-repo metrics.
	err = nt.ValidateMetrics(nomostest.SyncMetricsToLatestCommit(nt), func() error {
		err := nt.ValidateMultiRepoMetrics(nomostest.DefaultRootReconcilerName,
			nt.DefaultRootSyncObjectCount(),
			metrics.ResourceDeleted("Role"))
		if err != nil {
			return err
		}
		return nt.ValidateErrorMetricsNotFound()
	})
	if err != nil {
		nt.T.Error(err)
	}

	// Remove the lifecycle annotation from the namespace so that the namespace can be deleted after the test case.
	nt.RootRepos[configsync.RootSyncName].Add("acme/namespaces/shipping/ns.yaml", fake.NamespaceObject("shipping"))
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove the lifecycle annotation from Namespace")
	nt.WaitForRepoSyncs()
}

func TestPreventDeletionRole(t *testing.T) {
	nt := nomostest.New(t, nomostesting.Lifecycle)

	// Ensure the Namespace doesn't already exist.
	err := nt.ValidateNotFound("shipping-admin", "shipping", &rbacv1.Role{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Declare the Role with the lifecycle annotation, and ensure it is created.
	role := fake.RoleObject(core.Name("shipping-admin"), preventDeletion)
	role.Rules = []rbacv1.PolicyRule{{
		APIGroups: []string{""},
		Resources: []string{"configmaps"},
		Verbs:     []string{"get"},
	}}
	nt.RootRepos[configsync.RootSyncName].Add("acme/namespaces/shipping/ns.yaml", fake.NamespaceObject("shipping"))
	nt.RootRepos[configsync.RootSyncName].Add("acme/namespaces/shipping/role.yaml", role)
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("declare Role with prevent deletion lifecycle annotation")
	nt.WaitForRepoSyncs()

	err = nt.Validate("shipping-admin", "shipping", &rbacv1.Role{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Delete the declaration and ensure the Namespace isn't deleted.
	nt.RootRepos[configsync.RootSyncName].Remove("acme/namespaces/shipping/role.yaml")
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove Role declaration")
	nt.WaitForRepoSyncs()

	err = nt.Validate("shipping-admin", "shipping", &rbacv1.Role{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Remove the lifecycle annotation from the role so that the role can be deleted after the test case.
	delete(role.Annotations, common.LifecycleDeleteAnnotation)
	nt.RootRepos[configsync.RootSyncName].Add("acme/namespaces/shipping/role.yaml", role)
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove the lifecycle annotation from Role")
	nt.WaitForRepoSyncs()

	err = nt.ValidateMetrics(nomostest.SyncMetricsToLatestCommit(nt), func() error {
		return nt.ValidateErrorMetricsNotFound()
	})
	if err != nil {
		nt.T.Error(err)
	}
}

func TestPreventDeletionClusterRole(t *testing.T) {
	nt := nomostest.New(t, nomostesting.Lifecycle)

	// Ensure the ClusterRole doesn't already exist.
	err := nt.ValidateNotFound("test-admin", "", &rbacv1.ClusterRole{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Declare the ClusterRole with the lifecycle annotation, and ensure it is created.
	clusterRole := fake.ClusterRoleObject(core.Name("test-admin"), preventDeletion)
	clusterRole.Rules = []rbacv1.PolicyRule{{
		APIGroups: []string{""},
		Resources: []string{"configmaps"},
		Verbs:     []string{"get"},
	}}
	nt.RootRepos[configsync.RootSyncName].Add("acme/cluster/cr.yaml", clusterRole)
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("declare ClusterRole with prevent deletion lifecycle annotation")
	nt.WaitForRepoSyncs()

	err = nt.Validate("test-admin", "", &rbacv1.ClusterRole{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Delete the declaration and ensure the ClusterRole isn't deleted.
	nt.RootRepos[configsync.RootSyncName].Remove("acme/cluster/cr.yaml")
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove ClusterRole bar declaration")
	nt.WaitForRepoSyncs()

	err = nt.Validate("test-admin", "", &rbacv1.ClusterRole{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Remove the lifecycle annotation from the cluster-role so that it can be deleted after the test case.
	delete(clusterRole.Annotations, common.LifecycleDeleteAnnotation)
	nt.RootRepos[configsync.RootSyncName].Add("acme/cluster/cr.yaml", clusterRole)
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove the lifecycle annotation from ClusterRole")
	nt.WaitForRepoSyncs()

	err = nt.ValidateMetrics(nomostest.SyncMetricsToLatestCommit(nt), func() error {
		return nt.ValidateErrorMetricsNotFound()
	})
	if err != nil {
		nt.T.Error(err)
	}
}

func TestPreventDeletionImplicitNamespace(t *testing.T) {
	nt := nomostest.New(t, nomostesting.Lifecycle, ntopts.Unstructured, ntopts.SkipMultiRepo)

	const implicitNamespace = "delivery"

	role := fake.RoleObject(core.Name("configmap-getter"), core.Namespace(implicitNamespace))
	role.Rules = []rbacv1.PolicyRule{{
		APIGroups: []string{""},
		Resources: []string{"configmaps"},
		Verbs:     []string{"get"},
	}}
	nt.RootRepos[configsync.RootSyncName].Add("acme/role.yaml", role)
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("Declare configmap-getter Role")
	nt.WaitForRepoSyncs()

	err := nt.Validate(implicitNamespace, "", &corev1.Namespace{},
		nomostest.HasAnnotation(common.LifecycleDeleteAnnotation, common.PreventDeletion))
	if err != nil {
		nt.T.Fatal(err)
	}
	err = nt.Validate("configmap-getter", implicitNamespace, &rbacv1.Role{})
	if err != nil {
		nt.T.Fatal(err)
	}

	nt.RootRepos[configsync.RootSyncName].Remove("acme/role.yaml")
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("Remove configmap-getter Role")
	nt.WaitForRepoSyncs()

	// Ensure the Namespace wasn't deleted.
	err = nt.Validate(implicitNamespace, "", &corev1.Namespace{},
		nomostest.NotPendingDeletion)
	if err != nil {
		nt.T.Fatal(err)
	}

	// Remove the lifecycle annotation from the implicit namespace so that it can be deleted after the test case.
	nt.RootRepos[configsync.RootSyncName].Add("acme/ns.yaml", fake.NamespaceObject(implicitNamespace))
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove the lifecycle annotation from the implicit namespace")
	nt.WaitForRepoSyncs()
}

func skipAutopilotManagedNamespace(nt *nomostest.NT, ns string) bool {
	managedNS, found := util.AutopilotManagedNamespaces[ns]
	return found && managedNS && nt.IsGKEAutopilot
}

func TestPreventDeletionSpecialNamespaces(t *testing.T) {
	nt := nomostest.New(t, nomostesting.Lifecycle, ntopts.Unstructured)

	for ns := range differ.SpecialNamespaces {
		if !skipAutopilotManagedNamespace(nt, ns) {
			nt.RootRepos[configsync.RootSyncName].Add(fmt.Sprintf("acme/ns-%s.yaml", ns), fake.NamespaceObject(ns))
		}
	}
	nt.RootRepos[configsync.RootSyncName].Add("acme/ns-bookstore.yaml", fake.NamespaceObject("bookstore"))
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("Add special namespaces and one non-special namespace")
	nt.WaitForRepoSyncs()

	// Verify that the special namespaces have the `client.lifecycle.config.k8s.io/deletion: detach` annotation.
	for ns := range differ.SpecialNamespaces {
		if !skipAutopilotManagedNamespace(nt, ns) {
			if err := nt.Validate(ns, "", &corev1.Namespace{}, nomostest.HasAnnotation(common.LifecycleDeleteAnnotation, common.PreventDeletion)); err != nil {
				nt.T.Fatal(err)
			}
		}
	}

	// Verify that the bookstore namespace does not have the `client.lifecycle.config.k8s.io/deletion: detach` annotation.
	err := nt.Validate("bookstore", "", &corev1.Namespace{}, nomostest.MissingAnnotation(common.LifecycleDeleteAnnotation))
	if err != nil {
		nt.T.Fatal(err)
	}

	for ns := range differ.SpecialNamespaces {
		if !skipAutopilotManagedNamespace(nt, ns) {
			nt.RootRepos[configsync.RootSyncName].Remove(fmt.Sprintf("acme/ns-%s.yaml", ns))
		}
	}
	nt.RootRepos[configsync.RootSyncName].Remove("acme/ns-bookstore.yaml")
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("Remove namespaces")
	nt.WaitForRepoSyncs()

	// Verify that the special namespaces still exist and have the `client.lifecycle.config.k8s.io/deletion: detach` annotation.
	for ns := range differ.SpecialNamespaces {
		if !skipAutopilotManagedNamespace(nt, ns) {
			if err := nt.Validate(ns, "", &corev1.Namespace{}, nomostest.HasAnnotation(common.LifecycleDeleteAnnotation, common.PreventDeletion)); err != nil {
				nt.T.Fatal(err)
			}
		}
	}

	// Verify that the bookstore namespace is removed.
	// Use `nomostest.Retry` here because sometimes some resources have not been applied/pruned successfully
	// when Config Sync reports that a commit is synced successfully. go/cs-sync-status-accuracy proposes a
	// solution to fix this.
	_, err = nomostest.Retry(30*time.Second, func() error {
		return nt.ValidateNotFound("bookstore", "", &corev1.Namespace{})
	})
	if err != nil {
		nt.T.Fatal(err)
	}
}
<commit_msg>Added additional annotation validation to TestPreventDeletionRole (#125)<commit_after>// Copyright 2022 Google LLC
//
// 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 e2e

import (
	"fmt"
	"testing"
	"time"

	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	"kpt.dev/configsync/e2e/nomostest"
	"kpt.dev/configsync/e2e/nomostest/metrics"
	"kpt.dev/configsync/e2e/nomostest/ntopts"
	nomostesting "kpt.dev/configsync/e2e/nomostest/testing"
	"kpt.dev/configsync/pkg/api/configsync"
	"kpt.dev/configsync/pkg/core"
	"kpt.dev/configsync/pkg/syncer/differ"
	"kpt.dev/configsync/pkg/testing/fake"
	"kpt.dev/configsync/pkg/util"
	"sigs.k8s.io/cli-utils/pkg/common"
)

var preventDeletion = core.Annotation(common.LifecycleDeleteAnnotation, common.PreventDeletion)

func TestPreventDeletionNamespace(t *testing.T) {
	nt := nomostest.New(t, nomostesting.Lifecycle)

	// Ensure the Namespace doesn't already exist.
	err := nt.ValidateNotFound("shipping", "", &corev1.Namespace{})
	if err != nil {
		nt.T.Fatal(err)
	}

	role := fake.RoleObject(core.Name("shipping-admin"))
	role.Rules = []rbacv1.PolicyRule{{
		APIGroups: []string{""},
		Resources: []string{"configmaps"},
		Verbs:     []string{"get"},
	}}

	// Declare the Namespace with the lifecycle annotation, and ensure it is created.
	nt.RootRepos[configsync.RootSyncName].Add("acme/namespaces/shipping/ns.yaml",
		fake.NamespaceObject("shipping", preventDeletion))
	nt.RootRepos[configsync.RootSyncName].Add("acme/namespaces/shipping/role.yaml", role)
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("declare Namespace with prevent deletion lifecycle annotation")
	nt.WaitForRepoSyncs()

	err = nt.Validate("shipping", "", &corev1.Namespace{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Validate multi-repo metrics.
	err = nt.ValidateMetrics(nomostest.SyncMetricsToLatestCommit(nt), func() error {
		err := nt.ValidateMultiRepoMetrics(nomostest.DefaultRootReconcilerName,
			nt.DefaultRootSyncObjectCount()+2, // 2 for the test Namespace & Role
			metrics.ResourceCreated("Namespace"), metrics.ResourceCreated("Role"))
		if err != nil {
			return err
		}
		return nt.ValidateErrorMetricsNotFound()
	})
	if err != nil {
		nt.T.Error(err)
	}

	// Delete the declaration and ensure the Namespace isn't deleted.
	nt.RootRepos[configsync.RootSyncName].Remove("acme/namespaces/shipping/ns.yaml")
	nt.RootRepos[configsync.RootSyncName].Remove("acme/namespaces/shipping/role.yaml")
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove Namespace shipping declaration")
	nt.WaitForRepoSyncs()

	// Ensure we kept the undeclared Namespace that had the "deletion: prevent" annotation.
	err = nt.Validate("shipping", "", &corev1.Namespace{},
		nomostest.NotPendingDeletion)
	if err != nil {
		nt.T.Fatal(err)
	}
	// Ensure we deleted the undeclared Role that doesn't have the annotation.
	err = nt.ValidateNotFound("shipping-admin", "shipping", &rbacv1.Role{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Validate multi-repo metrics.
	err = nt.ValidateMetrics(nomostest.SyncMetricsToLatestCommit(nt), func() error {
		err := nt.ValidateMultiRepoMetrics(nomostest.DefaultRootReconcilerName,
			nt.DefaultRootSyncObjectCount(),
			metrics.ResourceDeleted("Role"))
		if err != nil {
			return err
		}
		return nt.ValidateErrorMetricsNotFound()
	})
	if err != nil {
		nt.T.Error(err)
	}

	// Remove the lifecycle annotation from the namespace so that the namespace can be deleted after the test case.
	nt.RootRepos[configsync.RootSyncName].Add("acme/namespaces/shipping/ns.yaml", fake.NamespaceObject("shipping"))
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove the lifecycle annotation from Namespace")
	nt.WaitForRepoSyncs()
}

func TestPreventDeletionRole(t *testing.T) {
	nt := nomostest.New(t, nomostesting.Lifecycle)

	// Ensure the Namespace doesn't already exist.
	err := nt.ValidateNotFound("shipping-admin", "shipping", &rbacv1.Role{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Declare the Role with the lifecycle annotation, and ensure it is created.
	role := fake.RoleObject(core.Name("shipping-admin"), preventDeletion)
	role.Rules = []rbacv1.PolicyRule{{
		APIGroups: []string{""},
		Resources: []string{"configmaps"},
		Verbs:     []string{"get"},
	}}
	nt.RootRepos[configsync.RootSyncName].Add("acme/namespaces/shipping/ns.yaml", fake.NamespaceObject("shipping"))
	nt.RootRepos[configsync.RootSyncName].Add("acme/namespaces/shipping/role.yaml", role)
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("declare Role with prevent deletion lifecycle annotation")
	nt.WaitForRepoSyncs()

	// ensure that the Role is created with the preventDeletion annotation
	err = nt.Validate("shipping-admin", "shipping", &rbacv1.Role{}, nomostest.HasAnnotation(common.LifecycleDeleteAnnotation, common.PreventDeletion))
	if err != nil {
		nt.T.Fatal(err)
	}

	// Delete the declaration and ensure the Namespace isn't deleted.
	nt.RootRepos[configsync.RootSyncName].Remove("acme/namespaces/shipping/role.yaml")
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove Role declaration")
	nt.WaitForRepoSyncs()

	err = nt.Validate("shipping-admin", "shipping", &rbacv1.Role{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Remove the lifecycle annotation from the role so that the role can be deleted after the test case.
	delete(role.Annotations, common.LifecycleDeleteAnnotation)
	nt.RootRepos[configsync.RootSyncName].Add("acme/namespaces/shipping/role.yaml", role)
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove the lifecycle annotation from Role")
	nt.WaitForRepoSyncs()

	err = nt.ValidateMetrics(nomostest.SyncMetricsToLatestCommit(nt), func() error {
		return nt.ValidateErrorMetricsNotFound()
	})
	if err != nil {
		nt.T.Error(err)
	}
}

func TestPreventDeletionClusterRole(t *testing.T) {
	nt := nomostest.New(t, nomostesting.Lifecycle)

	// Ensure the ClusterRole doesn't already exist.
	err := nt.ValidateNotFound("test-admin", "", &rbacv1.ClusterRole{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Declare the ClusterRole with the lifecycle annotation, and ensure it is created.
	clusterRole := fake.ClusterRoleObject(core.Name("test-admin"), preventDeletion)
	clusterRole.Rules = []rbacv1.PolicyRule{{
		APIGroups: []string{""},
		Resources: []string{"configmaps"},
		Verbs:     []string{"get"},
	}}
	nt.RootRepos[configsync.RootSyncName].Add("acme/cluster/cr.yaml", clusterRole)
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("declare ClusterRole with prevent deletion lifecycle annotation")
	nt.WaitForRepoSyncs()

	err = nt.Validate("test-admin", "", &rbacv1.ClusterRole{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Delete the declaration and ensure the ClusterRole isn't deleted.
	nt.RootRepos[configsync.RootSyncName].Remove("acme/cluster/cr.yaml")
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove ClusterRole bar declaration")
	nt.WaitForRepoSyncs()

	err = nt.Validate("test-admin", "", &rbacv1.ClusterRole{})
	if err != nil {
		nt.T.Fatal(err)
	}

	// Remove the lifecycle annotation from the cluster-role so that it can be deleted after the test case.
	delete(clusterRole.Annotations, common.LifecycleDeleteAnnotation)
	nt.RootRepos[configsync.RootSyncName].Add("acme/cluster/cr.yaml", clusterRole)
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove the lifecycle annotation from ClusterRole")
	nt.WaitForRepoSyncs()

	err = nt.ValidateMetrics(nomostest.SyncMetricsToLatestCommit(nt), func() error {
		return nt.ValidateErrorMetricsNotFound()
	})
	if err != nil {
		nt.T.Error(err)
	}
}

func TestPreventDeletionImplicitNamespace(t *testing.T) {
	nt := nomostest.New(t, nomostesting.Lifecycle, ntopts.Unstructured, ntopts.SkipMultiRepo)

	const implicitNamespace = "delivery"

	role := fake.RoleObject(core.Name("configmap-getter"), core.Namespace(implicitNamespace))
	role.Rules = []rbacv1.PolicyRule{{
		APIGroups: []string{""},
		Resources: []string{"configmaps"},
		Verbs:     []string{"get"},
	}}
	nt.RootRepos[configsync.RootSyncName].Add("acme/role.yaml", role)
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("Declare configmap-getter Role")
	nt.WaitForRepoSyncs()

	err := nt.Validate(implicitNamespace, "", &corev1.Namespace{},
		nomostest.HasAnnotation(common.LifecycleDeleteAnnotation, common.PreventDeletion))
	if err != nil {
		nt.T.Fatal(err)
	}
	err = nt.Validate("configmap-getter", implicitNamespace, &rbacv1.Role{})
	if err != nil {
		nt.T.Fatal(err)
	}

	nt.RootRepos[configsync.RootSyncName].Remove("acme/role.yaml")
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("Remove configmap-getter Role")
	nt.WaitForRepoSyncs()

	// Ensure the Namespace wasn't deleted.
	err = nt.Validate(implicitNamespace, "", &corev1.Namespace{},
		nomostest.NotPendingDeletion)
	if err != nil {
		nt.T.Fatal(err)
	}

	// Remove the lifecycle annotation from the implicit namespace so that it can be deleted after the test case.
	nt.RootRepos[configsync.RootSyncName].Add("acme/ns.yaml", fake.NamespaceObject(implicitNamespace))
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("remove the lifecycle annotation from the implicit namespace")
	nt.WaitForRepoSyncs()
}

func skipAutopilotManagedNamespace(nt *nomostest.NT, ns string) bool {
	managedNS, found := util.AutopilotManagedNamespaces[ns]
	return found && managedNS && nt.IsGKEAutopilot
}

func TestPreventDeletionSpecialNamespaces(t *testing.T) {
	nt := nomostest.New(t, nomostesting.Lifecycle, ntopts.Unstructured)

	for ns := range differ.SpecialNamespaces {
		if !skipAutopilotManagedNamespace(nt, ns) {
			nt.RootRepos[configsync.RootSyncName].Add(fmt.Sprintf("acme/ns-%s.yaml", ns), fake.NamespaceObject(ns))
		}
	}
	nt.RootRepos[configsync.RootSyncName].Add("acme/ns-bookstore.yaml", fake.NamespaceObject("bookstore"))
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("Add special namespaces and one non-special namespace")
	nt.WaitForRepoSyncs()

	// Verify that the special namespaces have the `client.lifecycle.config.k8s.io/deletion: detach` annotation.
	for ns := range differ.SpecialNamespaces {
		if !skipAutopilotManagedNamespace(nt, ns) {
			if err := nt.Validate(ns, "", &corev1.Namespace{}, nomostest.HasAnnotation(common.LifecycleDeleteAnnotation, common.PreventDeletion)); err != nil {
				nt.T.Fatal(err)
			}
		}
	}

	// Verify that the bookstore namespace does not have the `client.lifecycle.config.k8s.io/deletion: detach` annotation.
	err := nt.Validate("bookstore", "", &corev1.Namespace{}, nomostest.MissingAnnotation(common.LifecycleDeleteAnnotation))
	if err != nil {
		nt.T.Fatal(err)
	}

	for ns := range differ.SpecialNamespaces {
		if !skipAutopilotManagedNamespace(nt, ns) {
			nt.RootRepos[configsync.RootSyncName].Remove(fmt.Sprintf("acme/ns-%s.yaml", ns))
		}
	}
	nt.RootRepos[configsync.RootSyncName].Remove("acme/ns-bookstore.yaml")
	nt.RootRepos[configsync.RootSyncName].CommitAndPush("Remove namespaces")
	nt.WaitForRepoSyncs()

	// Verify that the special namespaces still exist and have the `client.lifecycle.config.k8s.io/deletion: detach` annotation.
	for ns := range differ.SpecialNamespaces {
		if !skipAutopilotManagedNamespace(nt, ns) {
			if err := nt.Validate(ns, "", &corev1.Namespace{}, nomostest.HasAnnotation(common.LifecycleDeleteAnnotation, common.PreventDeletion)); err != nil {
				nt.T.Fatal(err)
			}
		}
	}

	// Verify that the bookstore namespace is removed.
	// Use `nomostest.Retry` here because sometimes some resources have not been applied/pruned successfully
	// when Config Sync reports that a commit is synced successfully. go/cs-sync-status-accuracy proposes a
	// solution to fix this.
	_, err = nomostest.Retry(30*time.Second, func() error {
		return nt.ValidateNotFound("bookstore", "", &corev1.Namespace{})
	})
	if err != nil {
		nt.T.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// The sym package provides a way to iterate over and change the symbols in Go
// source files.
package sym

import (
	"bytes"
	"code.google.com/p/rog-go/exp/go/ast"
	"code.google.com/p/rog-go/exp/go/parser"
	"code.google.com/p/rog-go/exp/go/printer"
	"code.google.com/p/rog-go/exp/go/token"
	"code.google.com/p/rog-go/exp/go/types"
	"fmt"
	"go/build"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"sync"
)

// Info holds information about an identifier.
type Info struct {
	Pos      token.Pos   // position of symbol.
	Expr     ast.Expr    // expression for symbol (*ast.Ident or *ast.SelectorExpr)
	Ident    *ast.Ident  // identifier in parse tree (changing ident.Name changes the parse tree)
	ExprType types.Type  // type of expression.
	ReferPos token.Pos   // position of referred-to symbol.
	ReferObj *ast.Object // object referred to.
	Local    bool        // whether referred-to object is function-local.
	Universe bool        // whether referred-to object is in universe.
}

// Context holds the context for IterateSyms.
type Context struct {
	pkgMutex     sync.Mutex
	pkgCache     map[string]*ast.Package
	importer     types.Importer
	ChangedFiles map[string]*ast.File

	// FileSet holds the fileset used when importing packages.
	FileSet *token.FileSet

	// Logf is used to print warning messages.
	// If it is nil, no warning messages will be printed.
	Logf func(pos token.Pos, f string, a ...interface{})
}

func NewContext() *Context {
	ctxt := &Context{
		pkgCache:     make(map[string]*ast.Package),
		FileSet:      token.NewFileSet(),
		ChangedFiles: make(map[string]*ast.File),
	}
	ctxt.importer = ctxt.importerFunc()
	return ctxt
}

// Import imports and parses the package with the given path.
// It returns nil if it fails.
func (ctxt *Context) Import(path string) *ast.Package {
	// TODO return error.
	return ctxt.importer(path)
}

func (ctxt *Context) importerFunc() types.Importer {
	return func(path string) *ast.Package {
		ctxt.pkgMutex.Lock()
		defer ctxt.pkgMutex.Unlock()
		if pkg := ctxt.pkgCache[path]; pkg != nil {
			return pkg
		}
		cwd, _ := os.Getwd()
		bpkg, err := build.Import(path, cwd, 0)
		if err != nil {
			ctxt.logf(token.NoPos, "cannot find %q: %v", path, err)
			return nil
		}
		var files []string
		files = append(files, bpkg.GoFiles...)
		files = append(files, bpkg.CgoFiles...)
		files = append(files, bpkg.TestGoFiles...)
		for i, f := range files {
			files[i] = filepath.Join(bpkg.Dir, f)
		}
		pkgs, err := parser.ParseFiles(ctxt.FileSet, files, parser.ParseComments)
		if len(pkgs) == 0 {
			ctxt.logf(token.NoPos, "cannot parse package %q: %v", path, err)
			return nil
		}
		delete(pkgs, "documentation")
		for _, pkg := range pkgs {
			if ctxt.pkgCache[path] == nil {
				ctxt.pkgCache[path] = pkg
			} else {
				ctxt.logf(token.NoPos, "unexpected extra package %q in %q", pkg.Name, path)
			}
		}
		return ctxt.pkgCache[path]
	}
}

func (ctxt *Context) logf(pos token.Pos, f string, a ...interface{}) {
	if ctxt.Logf == nil {
		return
	}
	ctxt.Logf(pos, f, a...)
}

// IterateSyms calls visitf for each identifier in the given file.  If
// visitf returns false, the iteration stops.  If visitf changes
// info.Ident.Name, the file is added to ctxt.ChangedFiles.
func (ctxt *Context) IterateSyms(f *ast.File, visitf func(info *Info) bool) {
	var visit astVisitor
	ok := true
	local := false // TODO set to true inside function body
	visit = func(n ast.Node) bool {
		if !ok {
			return false
		}
		switch n := n.(type) {
		case *ast.ImportSpec:
			// If the file imports a package to ".", abort
			// because we don't support that (yet).
			if n.Name != nil && n.Name.Name == "." {
				ctxt.logf(n.Pos(), "import to . not supported")
				ok = false
				return false
			}
			return true

		case *ast.FuncDecl:
			// add object for init functions
			if n.Recv == nil && n.Name.Name == "init" {
				n.Name.Obj = ast.NewObj(ast.Fun, "init")
			}
			e := ast.Expr(n.Name)
			if n.Recv != nil {
				// It's a method, so we need to synthesise a
				// selector expression so that visitExpr doesn't
				// just see a blank name.
				if len(n.Recv.List) != 1 {
					ctxt.logf(n.Pos(), "expected one receiver only!")
					return true
				}
				e = &ast.SelectorExpr{
					X:   n.Recv.List[0].Type,
					Sel: n.Name,
				}
			}
			ok = ctxt.visitExpr(f, e, false, visitf)
			local = true
			ast.Walk(visit, n.Type)
			if n.Body != nil {
				ast.Walk(visit, n.Body)
			}
			local = false
			return false

		case *ast.Ident:
			ok = ctxt.visitExpr(f, n, local, visitf)
			return false

		case *ast.KeyValueExpr:
			// don't try to resolve the key part of a key-value
			// because it might be a map key which doesn't
			// need resolving, and we can't tell without being
			// complicated with types.
			ast.Walk(visit, n.Value)
			return false

		case *ast.SelectorExpr:
			ast.Walk(visit, n.X)
			ok = ctxt.visitExpr(f, n, local, visitf)
			return false

		case *ast.File:
			for _, d := range n.Decls {
				ast.Walk(visit, d)
			}
			return false
		}

		return true
	}
	ast.Walk(visit, f)
}

func (ctxt *Context) filename(f *ast.File) string {
	return ctxt.FileSet.Position(f.Package).Filename
}

func (ctxt *Context) visitExpr(f *ast.File, e ast.Expr, local bool, visitf func(*Info) bool) bool {
	var info Info
	info.Expr = e
	switch e := e.(type) {
	case *ast.Ident:
		if e.Name == "_" {
			return true
		}
		info.Pos = e.Pos()
		info.Ident = e
	case *ast.SelectorExpr:
		info.Pos = e.Sel.Pos()
		info.Ident = e.Sel
	}
	obj, t := types.ExprType(e, ctxt.importer)
	if obj == nil {
		ctxt.logf(e.Pos(), "no object for %s", pretty(e))
		return true
	}
	info.ExprType = t
	info.ReferObj = obj
	if parser.Universe.Lookup(obj.Name) != obj {
		info.ReferPos = types.DeclPos(obj)
		if info.ReferPos == token.NoPos {
			ctxt.logf(e.Pos(), "no declaration for %s", pretty(e))
			return true
		}
	} else {
		info.Universe = true
	}
	info.Local = local
	oldName := info.Ident.Name
	more := visitf(&info)
	if info.Ident.Name != oldName {
		ctxt.ChangedFiles[ctxt.filename(f)] = f
	}
	return more
}

// WriteFiles writes the given files, formatted as with gofmt.
func (ctxt *Context) WriteFiles(files map[string]*ast.File) error {
	// TODO should we try to continue changing files even after an error?
	for _, f := range files {
		name := ctxt.filename(f)
		newSrc, err := ctxt.gofmtFile(f)
		if err != nil {
			return fmt.Errorf("cannot format %q: %v", name, err)
		}
		err = ioutil.WriteFile(name, newSrc, 0666)
		if err != nil {
			return fmt.Errorf("cannot write %q: %v", name, err)
		}
	}
	return nil
}

// litToString converts from a string literal to a regular string.
func litToString(lit *ast.BasicLit) (v string) {
	if lit.Kind != token.STRING {
		panic("expected string")
	}
	v, err := strconv.Unquote(string(lit.Value))
	if err != nil {
		panic("cannot unquote")
	}
	return v
}

type astVisitor func(n ast.Node) bool

func (f astVisitor) Visit(n ast.Node) ast.Visitor {
	if f(n) {
		return f
	}
	return nil
}

var emptyFileSet = token.NewFileSet()

func pretty(n ast.Node) string {
	var b bytes.Buffer
	printer.Fprint(&b, emptyFileSet, n)
	return b.String()
}

var printConfig = &printer.Config{
	Mode:     printer.TabIndent | printer.UseSpaces,
	Tabwidth: 8,
}

func (ctxt *Context) gofmtFile(f *ast.File) ([]byte, error) {
	var buf bytes.Buffer
	_, err := printConfig.Fprint(&buf, ctxt.FileSet, f)
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}
<commit_msg>exp/go/sym: suppress warning about init functions<commit_after>// The sym package provides a way to iterate over and change the symbols in Go
// source files.
package sym

import (
	"bytes"
	"code.google.com/p/rog-go/exp/go/ast"
	"code.google.com/p/rog-go/exp/go/parser"
	"code.google.com/p/rog-go/exp/go/printer"
	"code.google.com/p/rog-go/exp/go/token"
	"code.google.com/p/rog-go/exp/go/types"
	"fmt"
	"go/build"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"sync"
)

// Info holds information about an identifier.
type Info struct {
	Pos      token.Pos   // position of symbol.
	Expr     ast.Expr    // expression for symbol (*ast.Ident or *ast.SelectorExpr)
	Ident    *ast.Ident  // identifier in parse tree (changing ident.Name changes the parse tree)
	ExprType types.Type  // type of expression.
	ReferPos token.Pos   // position of referred-to symbol.
	ReferObj *ast.Object // object referred to.
	Local    bool        // whether referred-to object is function-local.
	Universe bool        // whether referred-to object is in universe.
}

// Context holds the context for IterateSyms.
type Context struct {
	pkgMutex     sync.Mutex
	pkgCache     map[string]*ast.Package
	importer     types.Importer
	ChangedFiles map[string]*ast.File

	// FileSet holds the fileset used when importing packages.
	FileSet *token.FileSet

	// Logf is used to print warning messages.
	// If it is nil, no warning messages will be printed.
	Logf func(pos token.Pos, f string, a ...interface{})
}

func NewContext() *Context {
	ctxt := &Context{
		pkgCache:     make(map[string]*ast.Package),
		FileSet:      token.NewFileSet(),
		ChangedFiles: make(map[string]*ast.File),
	}
	ctxt.importer = ctxt.importerFunc()
	return ctxt
}

// Import imports and parses the package with the given path.
// It returns nil if it fails.
func (ctxt *Context) Import(path string) *ast.Package {
	// TODO return error.
	return ctxt.importer(path)
}

func (ctxt *Context) importerFunc() types.Importer {
	return func(path string) *ast.Package {
		ctxt.pkgMutex.Lock()
		defer ctxt.pkgMutex.Unlock()
		if pkg := ctxt.pkgCache[path]; pkg != nil {
			return pkg
		}
		cwd, _ := os.Getwd()
		bpkg, err := build.Import(path, cwd, 0)
		if err != nil {
			ctxt.logf(token.NoPos, "cannot find %q: %v", path, err)
			return nil
		}
		var files []string
		files = append(files, bpkg.GoFiles...)
		files = append(files, bpkg.CgoFiles...)
		files = append(files, bpkg.TestGoFiles...)
		for i, f := range files {
			files[i] = filepath.Join(bpkg.Dir, f)
		}
		pkgs, err := parser.ParseFiles(ctxt.FileSet, files, parser.ParseComments)
		if len(pkgs) == 0 {
			ctxt.logf(token.NoPos, "cannot parse package %q: %v", path, err)
			return nil
		}
		delete(pkgs, "documentation")
		for _, pkg := range pkgs {
			if ctxt.pkgCache[path] == nil {
				ctxt.pkgCache[path] = pkg
			} else {
				ctxt.logf(token.NoPos, "unexpected extra package %q in %q", pkg.Name, path)
			}
		}
		return ctxt.pkgCache[path]
	}
}

func (ctxt *Context) logf(pos token.Pos, f string, a ...interface{}) {
	if ctxt.Logf == nil {
		return
	}
	ctxt.Logf(pos, f, a...)
}

// IterateSyms calls visitf for each identifier in the given file.  If
// visitf returns false, the iteration stops.  If visitf changes
// info.Ident.Name, the file is added to ctxt.ChangedFiles.
func (ctxt *Context) IterateSyms(f *ast.File, visitf func(info *Info) bool) {
	var visit astVisitor
	ok := true
	local := false // TODO set to true inside function body
	visit = func(n ast.Node) bool {
		if !ok {
			return false
		}
		switch n := n.(type) {
		case *ast.ImportSpec:
			// If the file imports a package to ".", abort
			// because we don't support that (yet).
			if n.Name != nil && n.Name.Name == "." {
				ctxt.logf(n.Pos(), "import to . not supported")
				ok = false
				return false
			}
			return true

		case *ast.FuncDecl:
			// add object for init functions
			if n.Recv == nil && n.Name.Name == "init" {
				n.Name.Obj = ast.NewObj(ast.Fun, "init")
			}
			e := ast.Expr(n.Name)
			if n.Recv != nil {
				// It's a method, so we need to synthesise a
				// selector expression so that visitExpr doesn't
				// just see a blank name.
				if len(n.Recv.List) != 1 {
					ctxt.logf(n.Pos(), "expected one receiver only!")
					return true
				}
				e = &ast.SelectorExpr{
					X:   n.Recv.List[0].Type,
					Sel: n.Name,
				}
			}
			ok = ctxt.visitExpr(f, e, false, visitf)
			local = true
			ast.Walk(visit, n.Type)
			if n.Body != nil {
				ast.Walk(visit, n.Body)
			}
			local = false
			return false

		case *ast.Ident:
			ok = ctxt.visitExpr(f, n, local, visitf)
			return false

		case *ast.KeyValueExpr:
			// don't try to resolve the key part of a key-value
			// because it might be a map key which doesn't
			// need resolving, and we can't tell without being
			// complicated with types.
			ast.Walk(visit, n.Value)
			return false

		case *ast.SelectorExpr:
			ast.Walk(visit, n.X)
			ok = ctxt.visitExpr(f, n, local, visitf)
			return false

		case *ast.File:
			for _, d := range n.Decls {
				ast.Walk(visit, d)
			}
			return false
		}

		return true
	}
	ast.Walk(visit, f)
}

func (ctxt *Context) filename(f *ast.File) string {
	return ctxt.FileSet.Position(f.Package).Filename
}

func (ctxt *Context) visitExpr(f *ast.File, e ast.Expr, local bool, visitf func(*Info) bool) bool {
	var info Info
	info.Expr = e
	switch e := e.(type) {
	case *ast.Ident:
		if e.Name == "_" {
			return true
		}
		info.Pos = e.Pos()
		info.Ident = e
	case *ast.SelectorExpr:
		info.Pos = e.Sel.Pos()
		info.Ident = e.Sel
	}
	obj, t := types.ExprType(e, ctxt.importer)
	if obj == nil {
		ctxt.logf(e.Pos(), "no object for %s", pretty(e))
		return true
	}
	info.ExprType = t
	info.ReferObj = obj
	if parser.Universe.Lookup(obj.Name) != obj {
		info.ReferPos = types.DeclPos(obj)
		if info.ReferPos == token.NoPos {
			name := pretty(e)
			if name != "init" {
				ctxt.logf(e.Pos(), "no declaration for %s", pretty(e))
			}
			return true
		}
	} else {
		info.Universe = true
	}
	info.Local = local
	oldName := info.Ident.Name
	more := visitf(&info)
	if info.Ident.Name != oldName {
		ctxt.ChangedFiles[ctxt.filename(f)] = f
	}
	return more
}

// WriteFiles writes the given files, formatted as with gofmt.
func (ctxt *Context) WriteFiles(files map[string]*ast.File) error {
	// TODO should we try to continue changing files even after an error?
	for _, f := range files {
		name := ctxt.filename(f)
		newSrc, err := ctxt.gofmtFile(f)
		if err != nil {
			return fmt.Errorf("cannot format %q: %v", name, err)
		}
		err = ioutil.WriteFile(name, newSrc, 0666)
		if err != nil {
			return fmt.Errorf("cannot write %q: %v", name, err)
		}
	}
	return nil
}

// litToString converts from a string literal to a regular string.
func litToString(lit *ast.BasicLit) (v string) {
	if lit.Kind != token.STRING {
		panic("expected string")
	}
	v, err := strconv.Unquote(string(lit.Value))
	if err != nil {
		panic("cannot unquote")
	}
	return v
}

type astVisitor func(n ast.Node) bool

func (f astVisitor) Visit(n ast.Node) ast.Visitor {
	if f(n) {
		return f
	}
	return nil
}

var emptyFileSet = token.NewFileSet()

func pretty(n ast.Node) string {
	var b bytes.Buffer
	printer.Fprint(&b, emptyFileSet, n)
	return b.String()
}

var printConfig = &printer.Config{
	Mode:     printer.TabIndent | printer.UseSpaces,
	Tabwidth: 8,
}

func (ctxt *Context) gofmtFile(f *ast.File) ([]byte, error) {
	var buf bytes.Buffer
	_, err := printConfig.Fprint(&buf, ctxt.FileSet, f)
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}
<|endoftext|>


<commit_before>package shell

import (
	"bytes"
	"io"
	"testing"
)

func TestUnixReader_impl(t *testing.T) {
	var raw interface{}
	raw = new(UnixReader)
	if _, ok := raw.(io.Reader); !ok {
		t.Fatal("should be reader")
	}
}

func TestUnixReader(t *testing.T) {
	input := "one\r\ntwo\nthree\r\n"
	expected := "one\ntwo\nthree\n"

	r := &UnixReader{
		Reader: bytes.NewReader([]byte(input)),
	}

	result := new(bytes.Buffer)
	if _, err := io.Copy(result, r); err != nil {
		t.Fatalf("err: %s", err)
	}

	if result.String() != expected {
		t.Fatalf("bad: %#v", result.String())
	}
}

func TestUnixReader_unixOnly(t *testing.T) {
	input := "one\n\ntwo\nthree\n"
	expected := "one\n\ntwo\nthree\n"

	r := &UnixReader{
		Reader: bytes.NewReader([]byte(input)),
	}

	result := new(bytes.Buffer)
	if _, err := io.Copy(result, r); err != nil {
		t.Fatalf("err: %s", err)
	}

	if result.String() != expected {
		t.Fatalf("bad: %#v", result.String())
	}
}
<commit_msg>provisioner/shell: add more tests for UnixReader<commit_after>package shell

import (
	"bytes"
	"io"
	"testing"
)

func TestUnixReader_impl(t *testing.T) {
	var raw interface{}
	raw = new(UnixReader)
	if _, ok := raw.(io.Reader); !ok {
		t.Fatal("should be reader")
	}
}

func TestUnixReader(t *testing.T) {
	input := "one\r\ntwo\n\r\nthree\r\n"
	expected := "one\ntwo\n\nthree\n"

	r := &UnixReader{
		Reader: bytes.NewReader([]byte(input)),
	}

	result := new(bytes.Buffer)
	if _, err := io.Copy(result, r); err != nil {
		t.Fatalf("err: %s", err)
	}

	if result.String() != expected {
		t.Fatalf("bad: %#v", result.String())
	}
}

func TestUnixReader_unixOnly(t *testing.T) {
	input := "\none\n\ntwo\nthree\n\n"
	expected := "\none\n\ntwo\nthree\n\n"

	r := &UnixReader{
		Reader: bytes.NewReader([]byte(input)),
	}

	result := new(bytes.Buffer)
	if _, err := io.Copy(result, r); err != nil {
		t.Fatalf("err: %s", err)
	}

	if result.String() != expected {
		t.Fatalf("bad: %#v", result.String())
	}
}
<|endoftext|>