<commit_before>package main

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/RackSec/srslog"
	"github.com/buger/jsonparser"
	"github.com/urfave/cli"
)

var papertrailprefix string

func init() {
	app.Commands = append(app.Commands,
		cli.Command{
			Name:  "papertrail",
			Usage: "forward logs to papertrail",
			Action: func(c *cli.Context) {

				papertrailprefix = c.String("prefix")
				consumer := NewConsumer(globalFlags.Brokers, globalFlags.Groupid)
				consumer.Init(globalFlags.Topic)

				log.Println(c.String("papertrail"))

				go listenforchecks()
				papertrail := make(chan *data, 20)
				defer close(papertrail)
				go func(messages <-chan *Message, papertrail chan<- *data) {
					for msg := range messages {
						jsondata := &data{start: time.Now(), offset: msg.Offset}
						jsondata.Ts, _ = jsonparser.GetUnsafeString(msg.Data, "time")
						jsondata.Msg, _ = jsonparser.GetString(msg.Data, "log")
						jsondata.Level, _ = jsonparser.GetUnsafeString(msg.Data, "level")
						jsondata.Stream, _ = jsonparser.GetUnsafeString(msg.Data, "stream")
						jsondata.Service, _ = jsonparser.GetUnsafeString(msg.Data, "kubernetes", "labels", "app")
						jsondata.Host, _ = jsonparser.GetUnsafeString(msg.Data, "kubernetes", "host")
						jsondata.IPAddress, _ = jsonparser.GetUnsafeString(msg.Data, "ip_address")
						jsondata.ServerIP, _ = jsonparser.GetUnsafeString(msg.Data, "server_ip")
						jsondata.DockerImage, _ = jsonparser.GetUnsafeString(msg.Data, "docker_image")
						jsondata.ContainerName, _ = jsonparser.GetUnsafeString(msg.Data, "kubernetes", "pod_name")
						jsondata.ContainerID, _ = jsonparser.GetUnsafeString(msg.Data, "docker", "container_id")

						// NOTE: GetString does some allocations, which might cause some overhead.
						jsondata.Msg = strings.TrimSpace(jsondata.Msg)
						papertrail <- jsondata
					}
				}(consumer.Chan, papertrail)
				go Sender(papertrail, c.String("papertrail"), c.String("cert"))

				consumer.StartConsumingTopic()

				consumer.Wait()
			},
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:   "papertrail",
					Usage:  "The papertrail address where the logs should be forwarded",
					EnvVar: "PAPERTRAIL_ADDRESS",
				},
				cli.StringFlag{
					Name:   "cert",
					Usage:  "Papertrail root certificate path",
					Value:  "./papertrail-bundle.pem",
					EnvVar: "PAPERTRAIL_CERT",
				},
				cli.StringFlag{
					Name:   "prefix",
					Usage:  "String that is prepended to the system name",
					EnvVar: "PAPERTRAIL_PREFIX",
				},
			},
		})
}

type data struct {
	Ts            string `json:"ts"`
	Msg           string `json:"msg"`
	Level         string `json:"level"`
	Stream        string `json:"stream"`
	Service       string `json:"service"`
	Host          string `json:"host"`
	IPAddress     string `json:"ip_address"` // legacy
	ServerIP      string `json:"server_ip"`
	DockerImage   string `json:"docker_image"`
	ContainerName string `json:"container_name"`
	ContainerID   string `json:"container_id"`

	start  time.Time
	offset int64
}

// Sender receives valid entries from chan 'c' and uploads them into papertrail over an encrypted TCP connection.
func Sender(c <-chan *data, address, cert string) {
	w, err := srslog.DialWithTLSCertPath("tcp+tls", address, srslog.LOG_INFO, "kafkatopapertrail", cert)
	if err != nil {
		log.Fatal(err)
	}

	var timechan chan time.Duration
	if globalFlags.Verbose {
		timechan = make(chan time.Duration, 10)
		defer close(timechan)
		go func(timechan chan time.Duration) {
			var (
				eventcount int
				counter    time.Duration
				ticker     = time.NewTicker(time.Second * 30)
			)

		loop:
			for {
				select {
				case duration, ok := <-timechan:
					if ok {
						counter += duration
						eventcount++
					} else {
						break loop
					}
				case <-ticker.C:
					if eventcount != 0 {
						counter = counter / time.Duration(eventcount)
					}
					log.Printf("Sent %d messages during the last 30 seconds, averaging %s per message\n", eventcount, counter)
					eventcount = 0
					counter = 0
				}
			}
		}(timechan)
	}

	w.SetFormatter(LogFormatter)
	for msg := range c {
		switch msg.Level {
		case "INFO":
			err = w.Info(msg.String())
		case "ALERT":
			err = w.Alert(msg.String())
		case "CRIT", "CRITICAL":
			err = w.Crit(msg.String())
		case "DEBUG":
			err = w.Debug(msg.String())
		case "EMERG":
			err = w.Emerg(msg.String())
		case "ERR", "ERROR":
			err = w.Err(msg.String())
		case "NOTICE":
			err = w.Notice(msg.String())
		case "WARNING", "WARN":
			err = w.Warning(msg.String())
		default:
			err = w.Info(msg.String())
			//      log.Printf("Unknown log level: %s\n", msg.Level)
		}
		if err != nil {
			log.Println(err)
		}
		if globalFlags.Verbose {
			timechan <- time.Since(msg.start)
		}
	}
}

func (d *data) String() string {
	ts, _ := strconv.ParseInt(d.Ts, 10, 64)
	timestamp := time.Unix(ts, 0).Format(time.RFC3339)
	return fmt.Sprintf("%s|%s%s|%s|%s %s", timestamp, papertrailprefix, d.Service, d.ContainerName, timestamp, d.Msg)
}

func parseImage(image string) (repository, tag string) {
	// Docker image format is:
	// {registry}/(_|/r/{user_or_org})/{repository}:{tag}
	// Here we get the repository and tag from that format
	var (
		parts []string
	)
	parts = strings.Split(image, "/")
	image = parts[len(parts)-1]
	parts = strings.Split(image, ":")
	repository = parts[0]
	if len(parts) > 1 {
		tag = parts[1]
		if len(tag) > 8 {
			tag = tag[:8]
		}
	}
	return
}

// LogFormatter is a custom syslog formatter that uses log data to fill timestamp, hostname and tag instead of using local system information.
func LogFormatter(p srslog.Priority, hostname, tag, content string) (msg string) {
	parts := strings.SplitN(content, "|", 4)
	if len(parts) == 4 {
		msg = fmt.Sprintf("<%d>1 %s %s %s - - - %s", p, parts[0], parts[1], parts[2], parts[3])
	} else {
		msg = srslog.DefaultFormatter(p, hostname, tag, content)
	}
	return
}
<commit_msg>get prefix from environment<commit_after>package main

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/RackSec/srslog"
	"github.com/buger/jsonparser"
	"github.com/urfave/cli"
)

var papertrailprefix string

func init() {
	app.Commands = append(app.Commands,
		cli.Command{
			Name:  "papertrail",
			Usage: "forward logs to papertrail",
			Action: func(c *cli.Context) {

				papertrailprefix = c.String("prefix")
				consumer := NewConsumer(globalFlags.Brokers, globalFlags.Groupid)
				consumer.Init(globalFlags.Topic)

				log.Println(c.String("papertrail"))

				go listenforchecks()
				papertrail := make(chan *data, 20)
				defer close(papertrail)
				go func(messages <-chan *Message, papertrail chan<- *data) {
					for msg := range messages {
						jsondata := &data{start: time.Now(), offset: msg.Offset}
						jsondata.Ts, _ = jsonparser.GetUnsafeString(msg.Data, "time")
						jsondata.Msg, _ = jsonparser.GetString(msg.Data, "log")
						jsondata.Environment, _ = jsonparser.GetUnsafeString(msg.Data, "environment")
						jsondata.Level, _ = jsonparser.GetUnsafeString(msg.Data, "level")
						jsondata.Stream, _ = jsonparser.GetUnsafeString(msg.Data, "stream")
						jsondata.Service, _ = jsonparser.GetUnsafeString(msg.Data, "kubernetes", "labels", "app")
						jsondata.Host, _ = jsonparser.GetUnsafeString(msg.Data, "kubernetes", "host")
						jsondata.IPAddress, _ = jsonparser.GetUnsafeString(msg.Data, "ip_address")
						jsondata.ServerIP, _ = jsonparser.GetUnsafeString(msg.Data, "server_ip")
						jsondata.DockerImage, _ = jsonparser.GetUnsafeString(msg.Data, "docker_image")
						jsondata.ContainerName, _ = jsonparser.GetUnsafeString(msg.Data, "kubernetes", "pod_name")
						jsondata.ContainerID, _ = jsonparser.GetUnsafeString(msg.Data, "docker", "container_id")

						// NOTE: GetString does some allocations, which might cause some overhead.
						jsondata.Msg = strings.TrimSpace(jsondata.Msg)
						papertrail <- jsondata
					}
				}(consumer.Chan, papertrail)
				go Sender(papertrail, c.String("papertrail"), c.String("cert"))

				consumer.StartConsumingTopic()

				consumer.Wait()
			},
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:   "papertrail",
					Usage:  "The papertrail address where the logs should be forwarded",
					EnvVar: "PAPERTRAIL_ADDRESS",
				},
				cli.StringFlag{
					Name:   "cert",
					Usage:  "Papertrail root certificate path",
					Value:  "./papertrail-bundle.pem",
					EnvVar: "PAPERTRAIL_CERT",
				},
				cli.StringFlag{
					Name:   "prefix",
					Usage:  "String that is prepended to the system name",
					EnvVar: "PAPERTRAIL_PREFIX",
				},
			},
		})
}

type data struct {
	Ts            string `json:"ts"`
	Msg           string `json:"msg"`
	Level         string `json:"level"`
	Stream        string `json:"stream"`
	Service       string `json:"service"`
	Environment   string `json:"environment"`
	Host          string `json:"host"`
	IPAddress     string `json:"ip_address"` // legacy
	ServerIP      string `json:"server_ip"`
	DockerImage   string `json:"docker_image"`
	ContainerName string `json:"container_name"`
	ContainerID   string `json:"container_id"`

	start  time.Time
	offset int64
}

// Sender receives valid entries from chan 'c' and uploads them into papertrail over an encrypted TCP connection.
func Sender(c <-chan *data, address, cert string) {
	w, err := srslog.DialWithTLSCertPath("tcp+tls", address, srslog.LOG_INFO, "kafkatopapertrail", cert)
	if err != nil {
		log.Fatal(err)
	}

	var timechan chan time.Duration
	if globalFlags.Verbose {
		timechan = make(chan time.Duration, 10)
		defer close(timechan)
		go func(timechan chan time.Duration) {
			var (
				eventcount int
				counter    time.Duration
				ticker     = time.NewTicker(time.Second * 30)
			)

		loop:
			for {
				select {
				case duration, ok := <-timechan:
					if ok {
						counter += duration
						eventcount++
					} else {
						break loop
					}
				case <-ticker.C:
					if eventcount != 0 {
						counter = counter / time.Duration(eventcount)
					}
					log.Printf("Sent %d messages during the last 30 seconds, averaging %s per message\n", eventcount, counter)
					eventcount = 0
					counter = 0
				}
			}
		}(timechan)
	}

	w.SetFormatter(LogFormatter)
	for msg := range c {
		switch msg.Level {
		case "INFO":
			err = w.Info(msg.String())
		case "ALERT":
			err = w.Alert(msg.String())
		case "CRIT", "CRITICAL":
			err = w.Crit(msg.String())
		case "DEBUG":
			err = w.Debug(msg.String())
		case "EMERG":
			err = w.Emerg(msg.String())
		case "ERR", "ERROR":
			err = w.Err(msg.String())
		case "NOTICE":
			err = w.Notice(msg.String())
		case "WARNING", "WARN":
			err = w.Warning(msg.String())
		default:
			err = w.Info(msg.String())
			//      log.Printf("Unknown log level: %s\n", msg.Level)
		}
		if err != nil {
			log.Println(err)
		}
		if globalFlags.Verbose {
			timechan <- time.Since(msg.start)
		}
	}
}

func (d *data) String() string {
	prefix := papertrailprefix
	ts, _ := strconv.ParseInt(d.Ts, 10, 64)
	timestamp := time.Unix(ts, 0).Format(time.RFC3339)
	if d.Environment == "stg" {
		// A bit conflicting with the whole point of shared prefix but mostly a easy solution until some things are sorted.
		prefix = "staging-"
	}
	return fmt.Sprintf("%s|%s%s|%s|%s %s", timestamp, prefix, d.Service, d.ContainerName, timestamp, d.Msg)
}

func parseImage(image string) (repository, tag string) {
	// Docker image format is:
	// {registry}/(_|/r/{user_or_org})/{repository}:{tag}
	// Here we get the repository and tag from that format
	var (
		parts []string
	)
	parts = strings.Split(image, "/")
	image = parts[len(parts)-1]
	parts = strings.Split(image, ":")
	repository = parts[0]
	if len(parts) > 1 {
		tag = parts[1]
		if len(tag) > 8 {
			tag = tag[:8]
		}
	}
	return
}

// LogFormatter is a custom syslog formatter that uses log data to fill timestamp, hostname and tag instead of using local system information.
func LogFormatter(p srslog.Priority, hostname, tag, content string) (msg string) {
	parts := strings.SplitN(content, "|", 4)
	if len(parts) == 4 {
		msg = fmt.Sprintf("<%d>1 %s %s %s - - - %s", p, parts[0], parts[1], parts[2], parts[3])
	} else {
		msg = srslog.DefaultFormatter(p, hostname, tag, content)
	}
	return
}
<|endoftext|>


<commit_before>// Boringstreamer looks for mp3 files and broadcasts via http (live streaming.)
//
// $ boringstreamer
//
// or
//
// c:\>boringstreamer.exe
//
// recursively looks for mp3 files starting from current working directory and broadcasts on port 4444 for at most 42 concurrent http clients.
//
// See -h for details.
//
// Browse to listen (e.g. http://localhost:4444/)
package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/tcolgate/mp3"
)

var (
	addr           = flag.String("addr", ":4444", "listen on address (format: :port or host:port)")
	maxConnections = flag.Int("max", 42, "set maximum number of streaming connections")
	recursively    = flag.Bool("r", true, "recursively look for music starting from path")
	verbose        = flag.Bool("v", false, "display verbose messages")
)

// like /dev/null
type nullWriter struct {
}

func (nw nullWriter) Write(p []byte) (n int, err error) {
	return len(p), nil
}

type streamFrame []byte

// client's event
type broadcastResult struct {
	qid int
	err error
}

// After a start() mux broadcasts audio stream to it's listener clients.
// Clients subscribe() and unsubscribe by writing to result chanel.
type mux struct {
	sync.Mutex

	clients map[int]chan streamFrame // set of listener clients to be notified
	result  chan broadcastResult     // clients share broadcast success-failure here

	nextFile   chan string      // next file to be broadcast
	nextStream chan io.Reader   // next (ID3 stripped) raw audio stream
	nextFrame  chan streamFrame // next audio frame
}

// subscribe(ch) adds ch to the set of channels to be received on by the clients when a new audio frame is available.
// Returns uniq client id (qid) for ch and a broadcast result channel for the client.
// Returns -1, nil if too many clients are already listening.
// clients: qid, br := m.subscribe(ch)
func (m *mux) subscribe(ch chan streamFrame) (int, chan broadcastResult) {
	m.Lock()
	defer m.Unlock()
	// search for available qid
	qid := 0
	_, ok := m.clients[qid]
	for ; ok; _, ok = m.clients[qid] {
		if qid >= *maxConnections-1 {
			return -1, nil
		}
		qid++
	}
	m.clients[qid] = ch
	if *verbose {
		log.Printf("New connection (qid: %v), streaming to %v connections.", qid, len(m.clients))
	}

	return qid, m.result
}

// stripID3Header(r) reads file from r, strips id3v2 headers and returns the rest
// id3v2 tag details: id3.org
func stripID3Header(r io.Reader) io.Reader {
	buf, err := ioutil.ReadAll(r)
	if err != nil {
		log.Printf("Error: skipping file, stripID3Header(), err=%v", err)
		return bytes.NewReader(make([]byte, 0))
	}

	// TODO(fgergo) add ID3 v1 detection
	if string(buf[:3]) != "ID3" {
		return bytes.NewReader(buf) // no ID3 header
	}

	// The ID3v2 tag size is encoded in four bytes
	// where msb (bit 7) is set to zero in every byte,
	// ie. tag size is at most 2^28 (4*8-4=28).
	id3size := int32(buf[6])<<21 | int32(buf[7])<<14 | int32(buf[8])<<7 | int32(buf[9])
	id3size += 10 // calculated tag size is excluding the header => +10

	return bytes.NewReader(buf[id3size:])
}

// genFileList() periodically checks for files available from root and
// sends filenames down chan queue.
func genFileList(root string, queue chan string) {
	rand.Seed(time.Now().Unix()) // minimal randomness

	rescan := make(chan chan string)
	go func() {
		for {
			files := <-rescan
			filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
				if err != nil {
					return nil
				}
				if !info.Mode().IsRegular() {
					return nil
				}
				ok := strings.HasSuffix(strings.ToLower(info.Name()), ".mp3") // probably file is mp3
				if !info.IsDir() && !ok {
					return nil
				}
				files <- path // found file

				return nil
			})
			close(files)
			time.Sleep(1 * time.Second) // poll at least with 1Hz
		}
	}()

	// buffer and shuffle
	go func() {
		for {
			files := make(chan string)
			rescan <- files

			shuffled := make([]string, 0) // randomized set of files

			for f := range files {
				select {
				case queue <- f: // start playing as soon as possible
					if *verbose {
						fmt.Printf("Next: %v\n", f)
					}
					continue
				default:
					// shuffle files for random playback
					// (random permutation)
					if len(shuffled) == 0 {
						shuffled = append(shuffled, f)
					} else {
						i := rand.Intn(len(shuffled))
						shuffled = append(shuffled, shuffled[i])
						shuffled[i] = f
					}
				}
			}

			// queue shuffled files
			for _, f := range shuffled {
				queue <- f
				if *verbose {
					fmt.Printf("Next: %v\n", f)
				}
			}
		}
	}()
}

// start() initializes a multiplexer for raw audio streams
// e.g: m := new(mux).start(path)
func (m *mux) start(path string) *mux {
	m.result = make(chan broadcastResult)
	m.clients = make(map[int]chan streamFrame)

	m.nextFile = make(chan string)
	m.nextStream = make(chan io.Reader)
	m.nextFrame = make(chan streamFrame)

	// generate randomized list of files available from path
	genFileList(path, m.nextFile)

	// read file, strip ID3 header
	go func() {
		for {
			filename := <-m.nextFile
			f, err := os.Open(filename)
			if err != nil {
				log.Printf("Skipped \"%v\", err=%v", filename, err)
				continue
			}
			m.nextStream <- stripID3Header(f)
			if *verbose {
				fmt.Printf("Now playing: %v\n", filename)
			}
		}
	}()

	// decode stream to frames
	go func() {
		nullwriter := new(nullWriter)
		var cumwait time.Duration
		for {
			streamReader := <-m.nextStream
			d := mp3.NewDecoder(streamReader)
			var f mp3.Frame
			for {
				t0 := time.Now()
				tmp := log.Prefix()
				if !*verbose {
					log.SetOutput(nullwriter) // hack to silence mp3 debug/log output
				} else {
					log.SetPrefix("info: mp3 decode msg: ")
				}
				err := d.Decode(&f)
				log.SetPrefix(tmp)
				if !*verbose {
					log.SetOutput(os.Stderr)
				}
				if err == io.EOF {
					break
				}
				if err != nil {
					if *verbose {
						log.Printf("Skipping frame, d.Decode() err=%v", err)
					}
					continue
				}
				buf, err := ioutil.ReadAll(f.Reader())
				if err != nil {
					if *verbose {
						log.Printf("Skipping frame, ioutil.ReadAll() err=%v", err)
					}
					continue
				}
				m.nextFrame <- buf

				towait := f.Duration() - time.Now().Sub(t0)
				cumwait += towait // towait can be negative -> cumwait
				if cumwait > 4*time.Second {
					time.Sleep(cumwait)
					cumwait = 0
				}
			}
		}
	}()

	// broadcast frame to clients
	go func() {
		for {
			f := <-m.nextFrame
			// notify clients of new audio frame or let them quit
			for _, ch := range m.clients {
				ch <- f
				br := <-m.result // handle quitting clients
				if br.err != nil {
					m.Lock()
					close(m.clients[br.qid])
					delete(m.clients, br.qid)
					m.Unlock()
					if *verbose {
						log.Printf("Connection exited, qid: %v, error %v. Now streaming to %v connections.", br.qid, br.err, len(m.clients))
					}
				}
			}
		}
	}()

	return m
}

type streamHandler struct {
	stream *mux
}

// chrome and firefox play mp3 audio stream directly
func (sh streamHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	now := time.Now().UTC()
	frames := make(chan streamFrame)
	qid, br := sh.stream.subscribe(frames)
	if qid < 0 {
		log.Printf("Error: new connection request denied, already serving %v connections. See -h for details.", *maxConnections)
		w.WriteHeader(http.StatusTooManyRequests)
		return
	}

	w.WriteHeader(http.StatusOK)
	w.Header().Set("Date", now.Format(http.TimeFormat))
	w.Header().Set("Connection", "Keep-Alive")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Content-Type", "audio/mpeg")
	w.Header().Set("Server", "BoringStreamer/4.0")

	// browsers need ID3 tag to identify frames as media to be played
	// minimal id3 header to designate mp3 stream
	b := []byte{0x49, 0x44, 0x33, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
	_, err := io.Copy(w, bytes.NewReader(b))
	if err == nil {
		// broadcast mp3 stream to w
		broadcastTimeout := 4 * time.Second // timeout for slow clients
		result := make(chan error)
		for {
			buf := <-frames
			go func(r chan error, b []byte) {
				_, err = io.Copy(w, bytes.NewReader(b))
				if err == nil {
					w.(http.Flusher).Flush()
				}
				r <- err
			}(result, buf)
			select {
			case err = <-result:
				if err != nil {
					break
				}
				br <- broadcastResult{qid, nil} // frame streamed, no error, send ack
			case <-time.After(broadcastTimeout): // it's an error if io.Copy() is not finished within broadcastTimeout, ServeHTTP should exit
				err = errors.New(fmt.Sprintf("timeout: %v", broadcastTimeout))
			}
			if err != nil {
				break
			}
		}
	}
	br <- broadcastResult{qid, err} // error, send nack
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [flags] [path]\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "Browse to listen (e.g. http://localhost:4444/)\n\nflags:\n")
		flag.PrintDefaults()
	}
	flag.Parse()
	if len(flag.Args()) > 1 {
		flag.Usage()
		os.Exit(1)
	}

	path := "."
	switch len(flag.Args()) {
	case 0:
		if *verbose {
			fmt.Printf("Using path %#v, see -h for details.\n", path)
		}
	case 1:
		path = flag.Args()[0]
	}

	if *verbose {
		fmt.Printf("Looking for files available from \"%v\" ...\n", path)
	}

	// check if path is available
	matches, err := filepath.Glob(path)
	if err != nil || len(matches) != 1 {
		fmt.Fprintf(os.Stderr, "Error: \"%v\" unavailable.\n", path)
		os.Exit(1)
	}

	// initialize and start mp3 streamer
	http.Handle("/", streamHandler{new(mux).start(path)})
	if *verbose {
		fmt.Printf("Waiting for connections on %v\n", *addr)
	}

	err = http.ListenAndServe(*addr, nil)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>smaller waiting periods, quicker connections<commit_after>// Boringstreamer looks for mp3 files and broadcasts via http (live streaming.)
//
// $ boringstreamer
//
// or
//
// c:\>boringstreamer.exe
//
// recursively looks for mp3 files starting from current working directory and broadcasts on port 4444 for at most 42 concurrent http clients.
//
// See -h for details.
//
// Browse to listen (e.g. http://localhost:4444/)
package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/tcolgate/mp3"
)

var (
	addr           = flag.String("addr", ":4444", "listen on address (format: :port or host:port)")
	maxConnections = flag.Int("max", 42, "set maximum number of streaming connections")
	recursively    = flag.Bool("r", true, "recursively look for music starting from path")
	verbose        = flag.Bool("v", false, "display verbose messages")
)

// like /dev/null
type nullWriter struct {
}

func (nw nullWriter) Write(p []byte) (n int, err error) {
	return len(p), nil
}

type streamFrame []byte

// client's event
type broadcastResult struct {
	qid int
	err error
}

// After a start() mux broadcasts audio stream to it's listener clients.
// Clients subscribe() and unsubscribe by writing to result chanel.
type mux struct {
	sync.Mutex

	clients map[int]chan streamFrame // set of listener clients to be notified
	result  chan broadcastResult     // clients share broadcast success-failure here

	nextFile   chan string      // next file to be broadcast
	nextStream chan io.Reader   // next (ID3 stripped) raw audio stream
	nextFrame  chan streamFrame // next audio frame
}

// subscribe(ch) adds ch to the set of channels to be received on by the clients when a new audio frame is available.
// Returns uniq client id (qid) for ch and a broadcast result channel for the client.
// Returns -1, nil if too many clients are already listening.
// clients: qid, br := m.subscribe(ch)
func (m *mux) subscribe(ch chan streamFrame) (int, chan broadcastResult) {
	m.Lock()
	defer m.Unlock()
	// search for available qid
	qid := 0
	_, ok := m.clients[qid]
	for ; ok; _, ok = m.clients[qid] {
		if qid >= *maxConnections-1 {
			return -1, nil
		}
		qid++
	}
	m.clients[qid] = ch
	if *verbose {
		log.Printf("New connection (qid: %v), streaming to %v connections.", qid, len(m.clients))
	}

	return qid, m.result
}

// stripID3Header(r) reads file from r, strips id3v2 headers and returns the rest
// id3v2 tag details: id3.org
func stripID3Header(r io.Reader) io.Reader {
	buf, err := ioutil.ReadAll(r)
	if err != nil {
		log.Printf("Error: skipping file, stripID3Header(), err=%v", err)
		return bytes.NewReader(make([]byte, 0))
	}

	// TODO(fgergo) add ID3 v1 detection
	if string(buf[:3]) != "ID3" {
		return bytes.NewReader(buf) // no ID3 header
	}

	// The ID3v2 tag size is encoded in four bytes
	// where msb (bit 7) is set to zero in every byte,
	// ie. tag size is at most 2^28 (4*8-4=28).
	id3size := int32(buf[6])<<21 | int32(buf[7])<<14 | int32(buf[8])<<7 | int32(buf[9])
	id3size += 10 // calculated tag size is excluding the header => +10

	return bytes.NewReader(buf[id3size:])
}

// genFileList() periodically checks for files available from root and
// sends filenames down chan queue.
func genFileList(root string, queue chan string) {
	rand.Seed(time.Now().Unix()) // minimal randomness

	rescan := make(chan chan string)
	go func() {
		for {
			files := <-rescan
			filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
				if err != nil {
					return nil
				}
				if !info.Mode().IsRegular() {
					return nil
				}
				ok := strings.HasSuffix(strings.ToLower(info.Name()), ".mp3") // probably file is mp3
				if !info.IsDir() && !ok {
					return nil
				}
				files <- path // found file

				return nil
			})
			close(files)
			time.Sleep(1 * time.Second) // if no files are found, poll at least with 1Hz
		}
	}()

	// buffer and shuffle
	go func() {
		for {
			files := make(chan string)
			rescan <- files

			shuffled := make([]string, 0) // randomized set of files

			for f := range files {
				select {
				case queue <- f: // start playing as soon as possible
					if *verbose {
						fmt.Printf("Next: %v\n", f)
					}
					continue
				default:
					// shuffle files for random playback
					// (random permutation)
					if len(shuffled) == 0 {
						shuffled = append(shuffled, f)
					} else {
						i := rand.Intn(len(shuffled))
						shuffled = append(shuffled, shuffled[i])
						shuffled[i] = f
					}
				}
			}

			// queue shuffled files
			for _, f := range shuffled {
				queue <- f
				if *verbose {
					fmt.Printf("Next: %v\n", f)
				}
			}
		}
	}()
}

// start() initializes a multiplexer for raw audio streams
// e.g: m := new(mux).start(path)
func (m *mux) start(path string) *mux {
	m.result = make(chan broadcastResult)
	m.clients = make(map[int]chan streamFrame)

	m.nextFile = make(chan string)
	m.nextStream = make(chan io.Reader)
	m.nextFrame = make(chan streamFrame)

	// generate randomized list of files available from path
	genFileList(path, m.nextFile)

	// read file, strip ID3 header
	go func() {
		for {
			filename := <-m.nextFile
			f, err := os.Open(filename)
			if err != nil {
				log.Printf("Skipped \"%v\", err=%v", filename, err)
				continue
			}
			m.nextStream <- stripID3Header(f)
			if *verbose {
				fmt.Printf("Now playing: %v\n", filename)
			}
		}
	}()

	// decode stream to frames
	go func() {
		nullwriter := new(nullWriter)
		var cumwait time.Duration
		for {
			streamReader := <-m.nextStream
			d := mp3.NewDecoder(streamReader)
			var f mp3.Frame
			for {
				t0 := time.Now()
				tmp := log.Prefix()
				if !*verbose {
					log.SetOutput(nullwriter) // hack to silence mp3 debug/log output
				} else {
					log.SetPrefix("info: mp3 decode msg: ")
				}
				err := d.Decode(&f)
				log.SetPrefix(tmp)
				if !*verbose {
					log.SetOutput(os.Stderr)
				}
				if err == io.EOF {
					break
				}
				if err != nil {
					if *verbose {
						log.Printf("Skipping frame, d.Decode() err=%v", err)
					}
					continue
				}
				buf, err := ioutil.ReadAll(f.Reader())
				if err != nil {
					if *verbose {
						log.Printf("Skipping frame, ioutil.ReadAll() err=%v", err)
					}
					continue
				}
				m.nextFrame <- buf

				towait := f.Duration() - time.Now().Sub(t0)
				cumwait += towait // towait can be negative -> cumwait
				if cumwait > 1*time.Second {
					time.Sleep(cumwait)
					cumwait = 0
				}
			}
		}
	}()

	// broadcast frame to clients
	go func() {
		for {
			f := <-m.nextFrame
			// notify clients of new audio frame or let them quit
			for _, ch := range m.clients {
				ch <- f
				br := <-m.result // handle quitting clients
				if br.err != nil {
					m.Lock()
					close(m.clients[br.qid])
					delete(m.clients, br.qid)
					m.Unlock()
					if *verbose {
						log.Printf("Connection exited, qid: %v, error %v. Now streaming to %v connections.", br.qid, br.err, len(m.clients))
					}
				}
			}
		}
	}()

	return m
}

type streamHandler struct {
	stream *mux
}

// chrome and firefox play mp3 audio stream directly
func (sh streamHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	now := time.Now().UTC()
	frames := make(chan streamFrame)
	qid, br := sh.stream.subscribe(frames)
	if qid < 0 {
		log.Printf("Error: new connection request denied, already serving %v connections. See -h for details.", *maxConnections)
		w.WriteHeader(http.StatusTooManyRequests)
		return
	}

	w.WriteHeader(http.StatusOK)
	w.Header().Set("Date", now.Format(http.TimeFormat))
	w.Header().Set("Connection", "Keep-Alive")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Content-Type", "audio/mpeg")
	w.Header().Set("Server", "BoringStreamer/4.0")

	// browsers need ID3 tag to identify frames as media to be played
	// minimal id3 header to designate mp3 stream
	b := []byte{0x49, 0x44, 0x33, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
	_, err := io.Copy(w, bytes.NewReader(b))
	if err == nil {
		// broadcast mp3 stream to w
		broadcastTimeout := 4 * time.Second // timeout for slow clients
		result := make(chan error)
		for {
			buf := <-frames
			go func(r chan error, b []byte) {
				_, err = io.Copy(w, bytes.NewReader(b))
				if err == nil {
					w.(http.Flusher).Flush()
				}
				r <- err
			}(result, buf)
			select {
			case err = <-result:
				if err != nil {
					break
				}
				br <- broadcastResult{qid, nil} // frame streamed, no error, send ack
			case <-time.After(broadcastTimeout): // it's an error if io.Copy() is not finished within broadcastTimeout, ServeHTTP should exit
				err = errors.New(fmt.Sprintf("timeout: %v", broadcastTimeout))
			}
			if err != nil {
				break
			}
		}
	}
	br <- broadcastResult{qid, err} // error, send nack
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [flags] [path]\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "Browse to listen (e.g. http://localhost:4444/)\n\nflags:\n")
		flag.PrintDefaults()
	}
	flag.Parse()
	if len(flag.Args()) > 1 {
		flag.Usage()
		os.Exit(1)
	}

	path := "."
	switch len(flag.Args()) {
	case 0:
		if *verbose {
			fmt.Printf("Using path %#v, see -h for details.\n", path)
		}
	case 1:
		path = flag.Args()[0]
	}

	if *verbose {
		fmt.Printf("Looking for files available from \"%v\" ...\n", path)
	}

	// check if path is available
	matches, err := filepath.Glob(path)
	if err != nil || len(matches) != 1 {
		fmt.Fprintf(os.Stderr, "Error: \"%v\" unavailable.\n", path)
		os.Exit(1)
	}

	// initialize and start mp3 streamer
	http.Handle("/", streamHandler{new(mux).start(path)})
	if *verbose {
		fmt.Printf("Waiting for connections on %v\n", *addr)
	}

	err = http.ListenAndServe(*addr, nil)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package libstns

import (
	"github.com/BurntSushi/toml"
	"github.com/STNS/libnss_stns/settings"
)

type Config struct {
	ApiEndPoint     []string          `toml:"api_end_point"`
	RequestTimeOut  int               `toml:"request_timeout"`
	RequestRetry    int               `toml:"retry_request"`
	User            string            `toml:"user"`
	Password        string            `toml:"password"`
	SslVerify       bool              `toml:"ssl_verify"`
	WrapperCommand  string            `toml:"wrapper_path"`
	ChainSshWrapper string            `toml:"chain_ssh_wrapper"`
	HttpProxy       string            `toml:"http_proxy"`
	RequestHeader   map[string]string `toml:"request_header"`
	TlsCa           string            `toml:"tls_ca"`
	TlsCert         string            `toml:"tls_cert"`
	TlsKey          string            `toml:"tls_key"`
	UIDShift        int               `toml:"uid_shift"`
	GIDShift        int               `toml:"gid_shift"`
}

func LoadConfig(filePath string) (*Config, error) {
	var config Config

	defaultConfig(&config)
	_, err := toml.DecodeFile(filePath, &config)
	if err != nil {
		return nil, err
	}
	return &config, nil
}

func defaultConfig(config *Config) {
	config.RequestTimeOut = settings.HTTP_TIMEOUT
	config.RequestRetry = 1
	config.WrapperCommand = "/usr/local/bin/stns-query-wrapper"
	config.ApiEndPoint = []string{"http://localhost:1104"}
	config.UIDShift = 0
	config.GIDShift = 0
}
<commit_msg>request retry 3 times on default<commit_after>package libstns

import (
	"github.com/BurntSushi/toml"
	"github.com/STNS/libnss_stns/settings"
)

type Config struct {
	ApiEndPoint     []string          `toml:"api_end_point"`
	RequestTimeOut  int               `toml:"request_timeout"`
	RequestRetry    int               `toml:"retry_request"`
	User            string            `toml:"user"`
	Password        string            `toml:"password"`
	SslVerify       bool              `toml:"ssl_verify"`
	WrapperCommand  string            `toml:"wrapper_path"`
	ChainSshWrapper string            `toml:"chain_ssh_wrapper"`
	HttpProxy       string            `toml:"http_proxy"`
	RequestHeader   map[string]string `toml:"request_header"`
	TlsCa           string            `toml:"tls_ca"`
	TlsCert         string            `toml:"tls_cert"`
	TlsKey          string            `toml:"tls_key"`
	UIDShift        int               `toml:"uid_shift"`
	GIDShift        int               `toml:"gid_shift"`
}

func LoadConfig(filePath string) (*Config, error) {
	var config Config

	defaultConfig(&config)
	_, err := toml.DecodeFile(filePath, &config)
	if err != nil {
		return nil, err
	}
	return &config, nil
}

func defaultConfig(config *Config) {
	config.RequestTimeOut = settings.HTTP_TIMEOUT
	config.RequestRetry = 3
	config.WrapperCommand = "/usr/local/bin/stns-query-wrapper"
	config.ApiEndPoint = []string{"http://localhost:1104"}
	config.UIDShift = 0
	config.GIDShift = 0
}
<|endoftext|>


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

package fileutil

import (
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"

	"go.uber.org/zap"
)

const (
	// PrivateFileMode grants owner to read/write a file.
	PrivateFileMode = 0600
)

// IsDirWriteable checks if dir is writable by writing and removing a file
// to dir. It returns nil if dir is writable.
func IsDirWriteable(dir string) error {
	f, err := filepath.Abs(filepath.Join(dir, ".touch"))
	if err != nil {
		return err
	}
	if err := os.WriteFile(f, []byte(""), PrivateFileMode); err != nil {
		return err
	}
	return os.Remove(f)
}

// TouchDirAll is similar to os.MkdirAll. It creates directories with 0700 permission if any directory
// does not exists. TouchDirAll also ensures the given directory is writable.
func TouchDirAll(lg *zap.Logger, dir string) error {
	// If path is already a directory, MkdirAll does nothing and returns nil, so,
	// first check if dir exist with an expected permission mode.
	if Exist(dir) {
		err := CheckDirPermission(dir, PrivateDirMode)
		if err != nil {
			lg.Warn("check file permission", zap.Error(err))
		}
	} else {
		err := os.MkdirAll(dir, PrivateDirMode)
		if err != nil {
			// if mkdirAll("a/text") and "text" is not
			// a directory, this will return syscall.ENOTDIR
			return err
		}
	}

	return IsDirWriteable(dir)
}

// CreateDirAll is similar to TouchDirAll but returns error
// if the deepest directory was not empty.
func CreateDirAll(lg *zap.Logger, dir string) error {
	err := TouchDirAll(lg, dir)
	if err == nil {
		var ns []string
		ns, err = ReadDir(dir)
		if err != nil {
			return err
		}
		if len(ns) != 0 {
			err = fmt.Errorf("expected %q to be empty, got %q", dir, ns)
		}
	}
	return err
}

// Exist returns true if a file or directory exists.
func Exist(name string) bool {
	_, err := os.Stat(name)
	return err == nil
}

// DirEmpty returns true if a directory empty and can access.
func DirEmpty(name string) bool {
	ns, err := ReadDir(name)
	return len(ns) == 0 && err == nil
}

// ZeroToEnd zeros a file starting from SEEK_CUR to its SEEK_END. May temporarily
// shorten the length of the file.
func ZeroToEnd(f *os.File) error {
	// TODO: support FALLOC_FL_ZERO_RANGE
	off, err := f.Seek(0, io.SeekCurrent)
	if err != nil {
		return err
	}
	lenf, lerr := f.Seek(0, io.SeekEnd)
	if lerr != nil {
		return lerr
	}
	if err = f.Truncate(off); err != nil {
		return err
	}
	// make sure blocks remain allocated
	if err = Preallocate(f, lenf, true); err != nil {
		return err
	}
	_, err = f.Seek(off, io.SeekStart)
	return err
}

// CheckDirPermission checks permission on an existing dir.
// Returns error if dir is empty or exist with a different permission than specified.
func CheckDirPermission(dir string, perm os.FileMode) error {
	if !Exist(dir) {
		return fmt.Errorf("directory %q empty, cannot check permission", dir)
	}
	//check the existing permission on the directory
	dirInfo, err := os.Stat(dir)
	if err != nil {
		return err
	}
	dirMode := dirInfo.Mode().Perm()
	if dirMode != perm {
		err = fmt.Errorf("directory %q exist, but the permission is %q. The recommended permission is %q to prevent possible unprivileged access to the data", dir, dirInfo.Mode(), os.FileMode(PrivateDirMode))
		return err
	}
	return nil
}

// RemoveMatchFile deletes file if matchFunc is true on an existing dir
// Returns error if the dir does not exist or remove file fail
func RemoveMatchFile(lg *zap.Logger, dir string, matchFunc func(fileName string) bool) error {
	if lg == nil {
		lg = zap.NewNop()
	}
	if !Exist(dir) {
		return fmt.Errorf("directory %s does not exist", dir)
	}
	fileNames, err := ReadDir(dir)
	if err != nil {
		return err
	}
	var removeFailedFiles []string
	for _, fileName := range fileNames {
		if matchFunc(fileName) {
			file := filepath.Join(dir, fileName)
			if err = os.Remove(file); err != nil {
				removeFailedFiles = append(removeFailedFiles, fileName)
				lg.Error("remove file failed",
					zap.String("file", file),
					zap.Error(err))
				continue
			}
		}
	}
	if len(removeFailedFiles) != 0 {
		return fmt.Errorf("remove file(s) %v error", removeFailedFiles)
	}
	return nil
}

// ListFiles lists files if matchFunc is true on an existing dir
// Returns error if the dir does not exist
func ListFiles(dir string, matchFunc func(fileName string) bool) ([]string, error) {
	var files []string
	err := filepath.Walk(dir, func(path string, info fs.FileInfo, err error) error {
		if matchFunc(path) {
			files = append(files, path)
		}
		return nil
	})
	return files, err
}
<commit_msg>Fix syntax errors in comments<commit_after>// Copyright 2015 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package fileutil

import (
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"

	"go.uber.org/zap"
)

const (
	// PrivateFileMode grants owner to read/write a file.
	PrivateFileMode = 0600
)

// IsDirWriteable checks if dir is writable by writing and removing a file
// to dir. It returns nil if dir is writable.
func IsDirWriteable(dir string) error {
	f, err := filepath.Abs(filepath.Join(dir, ".touch"))
	if err != nil {
		return err
	}
	if err := os.WriteFile(f, []byte(""), PrivateFileMode); err != nil {
		return err
	}
	return os.Remove(f)
}

// TouchDirAll is similar to os.MkdirAll. It creates directories with 0700 permission if any directory
// does not exists. TouchDirAll also ensures the given directory is writable.
func TouchDirAll(lg *zap.Logger, dir string) error {
	// If path is already a directory, MkdirAll does nothing and returns nil, so,
	// first check if dir exists with an expected permission mode.
	if Exist(dir) {
		err := CheckDirPermission(dir, PrivateDirMode)
		if err != nil {
			lg.Warn("check file permission", zap.Error(err))
		}
	} else {
		err := os.MkdirAll(dir, PrivateDirMode)
		if err != nil {
			// if mkdirAll("a/text") and "text" is not
			// a directory, this will return syscall.ENOTDIR
			return err
		}
	}

	return IsDirWriteable(dir)
}

// CreateDirAll is similar to TouchDirAll but returns error
// if the deepest directory was not empty.
func CreateDirAll(lg *zap.Logger, dir string) error {
	err := TouchDirAll(lg, dir)
	if err == nil {
		var ns []string
		ns, err = ReadDir(dir)
		if err != nil {
			return err
		}
		if len(ns) != 0 {
			err = fmt.Errorf("expected %q to be empty, got %q", dir, ns)
		}
	}
	return err
}

// Exist returns true if a file or directory exists.
func Exist(name string) bool {
	_, err := os.Stat(name)
	return err == nil
}

// DirEmpty returns true if a directory empty and can access.
func DirEmpty(name string) bool {
	ns, err := ReadDir(name)
	return len(ns) == 0 && err == nil
}

// ZeroToEnd zeros a file starting from SEEK_CUR to its SEEK_END. May temporarily
// shorten the length of the file.
func ZeroToEnd(f *os.File) error {
	// TODO: support FALLOC_FL_ZERO_RANGE
	off, err := f.Seek(0, io.SeekCurrent)
	if err != nil {
		return err
	}
	lenf, lerr := f.Seek(0, io.SeekEnd)
	if lerr != nil {
		return lerr
	}
	if err = f.Truncate(off); err != nil {
		return err
	}
	// make sure blocks remain allocated
	if err = Preallocate(f, lenf, true); err != nil {
		return err
	}
	_, err = f.Seek(off, io.SeekStart)
	return err
}

// CheckDirPermission checks permission on an existing dir.
// Returns error if dir is empty or exist with a different permission than specified.
func CheckDirPermission(dir string, perm os.FileMode) error {
	if !Exist(dir) {
		return fmt.Errorf("directory %q empty, cannot check permission", dir)
	}
	//check the existing permission on the directory
	dirInfo, err := os.Stat(dir)
	if err != nil {
		return err
	}
	dirMode := dirInfo.Mode().Perm()
	if dirMode != perm {
		err = fmt.Errorf("directory %q exist, but the permission is %q. The recommended permission is %q to prevent possible unprivileged access to the data", dir, dirInfo.Mode(), os.FileMode(PrivateDirMode))
		return err
	}
	return nil
}

// RemoveMatchFile deletes file if matchFunc is true on an existing dir
// Returns error if the dir does not exist or remove file fail
func RemoveMatchFile(lg *zap.Logger, dir string, matchFunc func(fileName string) bool) error {
	if lg == nil {
		lg = zap.NewNop()
	}
	if !Exist(dir) {
		return fmt.Errorf("directory %s does not exist", dir)
	}
	fileNames, err := ReadDir(dir)
	if err != nil {
		return err
	}
	var removeFailedFiles []string
	for _, fileName := range fileNames {
		if matchFunc(fileName) {
			file := filepath.Join(dir, fileName)
			if err = os.Remove(file); err != nil {
				removeFailedFiles = append(removeFailedFiles, fileName)
				lg.Error("remove file failed",
					zap.String("file", file),
					zap.Error(err))
				continue
			}
		}
	}
	if len(removeFailedFiles) != 0 {
		return fmt.Errorf("remove file(s) %v error", removeFailedFiles)
	}
	return nil
}

// ListFiles lists files if matchFunc is true on an existing dir
// Returns error if the dir does not exist
func ListFiles(dir string, matchFunc func(fileName string) bool) ([]string, error) {
	var files []string
	err := filepath.Walk(dir, func(path string, info fs.FileInfo, err error) error {
		if matchFunc(path) {
			files = append(files, path)
		}
		return nil
	})
	return files, err
}
<|endoftext|>


<commit_before>package quic

import (
	"bytes"
	"errors"
	"io"

	"github.com/lucas-clemente/quic-go/protocol"
	"github.com/lucas-clemente/quic-go/qerr"
	"github.com/lucas-clemente/quic-go/utils"
)

var (
	errPacketNumberLenNotSet          = errors.New("PublicHeader: PacketNumberLen not set")
	errResetAndVersionFlagSet         = errors.New("PublicHeader: Reset Flag and Version Flag should not be set at the same time")
	errReceivedTruncatedConnectionID  = qerr.Error(qerr.InvalidPacketHeader, "receiving packets with truncated ConnectionID is not supported")
	errInvalidConnectionID            = qerr.Error(qerr.InvalidPacketHeader, "connection ID cannot be 0")
	errGetLengthOnlyForRegularPackets = errors.New("PublicHeader: GetLength can only be called for regular packets")
)

// The publicHeader of a QUIC packet
type publicHeader struct {
	Raw                  []byte
	VersionFlag          bool
	ResetFlag            bool
	ConnectionID         protocol.ConnectionID
	TruncateConnectionID bool
	VersionNumber        protocol.VersionNumber
	QuicVersion          uint32
	PacketNumberLen      protocol.PacketNumberLen
	PacketNumber         protocol.PacketNumber
	DiversificationNonce []byte
}

// WritePublicHeader writes a public header
func (h *publicHeader) WritePublicHeader(b *bytes.Buffer, version protocol.VersionNumber) error {
	publicFlagByte := uint8(0x00)
	if h.VersionFlag && h.ResetFlag {
		return errResetAndVersionFlagSet
	}
	if h.VersionFlag {
		publicFlagByte |= 0x01
	}
	if h.ResetFlag {
		publicFlagByte |= 0x02
	}
	if !h.TruncateConnectionID {
		if version < protocol.Version33 {
			publicFlagByte |= 0x0c
		} else {
			publicFlagByte |= 0x08
		}
	}
	if len(h.DiversificationNonce) > 0 {
		if len(h.DiversificationNonce) != 32 {
			return errors.New("invalid diversification nonce length")
		}
		publicFlagByte |= 0x04
	}

	if !h.ResetFlag && !h.VersionFlag {
		switch h.PacketNumberLen {
		case protocol.PacketNumberLen1:
			publicFlagByte |= 0x00
		case protocol.PacketNumberLen2:
			publicFlagByte |= 0x10
		case protocol.PacketNumberLen4:
			publicFlagByte |= 0x20
		case protocol.PacketNumberLen6:
			publicFlagByte |= 0x30
		}
	}

	b.WriteByte(publicFlagByte)

	if !h.TruncateConnectionID {
		utils.WriteUint64(b, uint64(h.ConnectionID))
	}

	if len(h.DiversificationNonce) > 0 {
		b.Write(h.DiversificationNonce)
	}

	if !h.ResetFlag && !h.VersionFlag {
		switch h.PacketNumberLen {
		case protocol.PacketNumberLen1:
			b.WriteByte(uint8(h.PacketNumber))
		case protocol.PacketNumberLen2:
			utils.WriteUint16(b, uint16(h.PacketNumber))
		case protocol.PacketNumberLen4:
			utils.WriteUint32(b, uint32(h.PacketNumber))
		case protocol.PacketNumberLen6:
			utils.WriteUint48(b, uint64(h.PacketNumber))
		default:
			return errPacketNumberLenNotSet
		}
	}

	return nil
}

// parsePublicHeader parses a QUIC packet's public header
func parsePublicHeader(b io.ByteReader) (*publicHeader, error) {
	header := &publicHeader{}

	// First byte
	publicFlagByte, err := b.ReadByte()
	if err != nil {
		return nil, err
	}
	header.VersionFlag = publicFlagByte&0x01 > 0
	header.ResetFlag = publicFlagByte&0x02 > 0

	// TODO: Add this check when we drop support for <v33
	// if publicFlagByte&0x04 > 0 {
	// 	return nil, errors.New("diversification nonces should only be sent by servers")
	// }

	if publicFlagByte&0x08 == 0 {
		return nil, errReceivedTruncatedConnectionID
	}

	switch publicFlagByte & 0x30 {
	case 0x30:
		header.PacketNumberLen = protocol.PacketNumberLen6
	case 0x20:
		header.PacketNumberLen = protocol.PacketNumberLen4
	case 0x10:
		header.PacketNumberLen = protocol.PacketNumberLen2
	case 0x00:
		header.PacketNumberLen = protocol.PacketNumberLen1
	}

	// Connection ID
	connID, err := utils.ReadUint64(b)
	if err != nil {
		return nil, err
	}
	header.ConnectionID = protocol.ConnectionID(connID)
	if header.ConnectionID == 0 {
		return nil, errInvalidConnectionID
	}

	// Version (optional)
	if header.VersionFlag {
		var versionTag uint32
		versionTag, err = utils.ReadUint32(b)
		if err != nil {
			return nil, err
		}
		header.VersionNumber = protocol.VersionTagToNumber(versionTag)
	}

	// Packet number
	packetNumber, err := utils.ReadUintN(b, uint8(header.PacketNumberLen))
	if err != nil {
		return nil, err
	}
	header.PacketNumber = protocol.PacketNumber(packetNumber)

	return header, nil
}

// GetLength gets the length of the publicHeader in bytes
// can only be called for regular packets
func (h *publicHeader) GetLength() (protocol.ByteCount, error) {
	if h.VersionFlag || h.ResetFlag {
		return 0, errGetLengthOnlyForRegularPackets
	}

	length := protocol.ByteCount(1) // 1 byte for public flags
	if h.PacketNumberLen != protocol.PacketNumberLen1 && h.PacketNumberLen != protocol.PacketNumberLen2 && h.PacketNumberLen != protocol.PacketNumberLen4 && h.PacketNumberLen != protocol.PacketNumberLen6 {
		return 0, errPacketNumberLenNotSet
	}
	if !h.TruncateConnectionID {
		length += 8 // 8 bytes for the connection ID
	}
	length += protocol.ByteCount(len(h.DiversificationNonce))
	length += protocol.ByteCount(h.PacketNumberLen)
	return length, nil
}
<commit_msg>remove unused member from PublicHeader<commit_after>package quic

import (
	"bytes"
	"errors"
	"io"

	"github.com/lucas-clemente/quic-go/protocol"
	"github.com/lucas-clemente/quic-go/qerr"
	"github.com/lucas-clemente/quic-go/utils"
)

var (
	errPacketNumberLenNotSet          = errors.New("PublicHeader: PacketNumberLen not set")
	errResetAndVersionFlagSet         = errors.New("PublicHeader: Reset Flag and Version Flag should not be set at the same time")
	errReceivedTruncatedConnectionID  = qerr.Error(qerr.InvalidPacketHeader, "receiving packets with truncated ConnectionID is not supported")
	errInvalidConnectionID            = qerr.Error(qerr.InvalidPacketHeader, "connection ID cannot be 0")
	errGetLengthOnlyForRegularPackets = errors.New("PublicHeader: GetLength can only be called for regular packets")
)

// The publicHeader of a QUIC packet
type publicHeader struct {
	Raw                  []byte
	VersionFlag          bool
	ResetFlag            bool
	ConnectionID         protocol.ConnectionID
	TruncateConnectionID bool
	VersionNumber        protocol.VersionNumber
	PacketNumberLen      protocol.PacketNumberLen
	PacketNumber         protocol.PacketNumber
	DiversificationNonce []byte
}

// WritePublicHeader writes a public header
func (h *publicHeader) WritePublicHeader(b *bytes.Buffer, version protocol.VersionNumber) error {
	publicFlagByte := uint8(0x00)
	if h.VersionFlag && h.ResetFlag {
		return errResetAndVersionFlagSet
	}
	if h.VersionFlag {
		publicFlagByte |= 0x01
	}
	if h.ResetFlag {
		publicFlagByte |= 0x02
	}
	if !h.TruncateConnectionID {
		if version < protocol.Version33 {
			publicFlagByte |= 0x0c
		} else {
			publicFlagByte |= 0x08
		}
	}
	if len(h.DiversificationNonce) > 0 {
		if len(h.DiversificationNonce) != 32 {
			return errors.New("invalid diversification nonce length")
		}
		publicFlagByte |= 0x04
	}

	if !h.ResetFlag && !h.VersionFlag {
		switch h.PacketNumberLen {
		case protocol.PacketNumberLen1:
			publicFlagByte |= 0x00
		case protocol.PacketNumberLen2:
			publicFlagByte |= 0x10
		case protocol.PacketNumberLen4:
			publicFlagByte |= 0x20
		case protocol.PacketNumberLen6:
			publicFlagByte |= 0x30
		}
	}

	b.WriteByte(publicFlagByte)

	if !h.TruncateConnectionID {
		utils.WriteUint64(b, uint64(h.ConnectionID))
	}

	if len(h.DiversificationNonce) > 0 {
		b.Write(h.DiversificationNonce)
	}

	if !h.ResetFlag && !h.VersionFlag {
		switch h.PacketNumberLen {
		case protocol.PacketNumberLen1:
			b.WriteByte(uint8(h.PacketNumber))
		case protocol.PacketNumberLen2:
			utils.WriteUint16(b, uint16(h.PacketNumber))
		case protocol.PacketNumberLen4:
			utils.WriteUint32(b, uint32(h.PacketNumber))
		case protocol.PacketNumberLen6:
			utils.WriteUint48(b, uint64(h.PacketNumber))
		default:
			return errPacketNumberLenNotSet
		}
	}

	return nil
}

// parsePublicHeader parses a QUIC packet's public header
func parsePublicHeader(b io.ByteReader) (*publicHeader, error) {
	header := &publicHeader{}

	// First byte
	publicFlagByte, err := b.ReadByte()
	if err != nil {
		return nil, err
	}
	header.VersionFlag = publicFlagByte&0x01 > 0
	header.ResetFlag = publicFlagByte&0x02 > 0

	// TODO: Add this check when we drop support for <v33
	// if publicFlagByte&0x04 > 0 {
	// 	return nil, errors.New("diversification nonces should only be sent by servers")
	// }

	if publicFlagByte&0x08 == 0 {
		return nil, errReceivedTruncatedConnectionID
	}

	switch publicFlagByte & 0x30 {
	case 0x30:
		header.PacketNumberLen = protocol.PacketNumberLen6
	case 0x20:
		header.PacketNumberLen = protocol.PacketNumberLen4
	case 0x10:
		header.PacketNumberLen = protocol.PacketNumberLen2
	case 0x00:
		header.PacketNumberLen = protocol.PacketNumberLen1
	}

	// Connection ID
	connID, err := utils.ReadUint64(b)
	if err != nil {
		return nil, err
	}
	header.ConnectionID = protocol.ConnectionID(connID)
	if header.ConnectionID == 0 {
		return nil, errInvalidConnectionID
	}

	// Version (optional)
	if header.VersionFlag {
		var versionTag uint32
		versionTag, err = utils.ReadUint32(b)
		if err != nil {
			return nil, err
		}
		header.VersionNumber = protocol.VersionTagToNumber(versionTag)
	}

	// Packet number
	packetNumber, err := utils.ReadUintN(b, uint8(header.PacketNumberLen))
	if err != nil {
		return nil, err
	}
	header.PacketNumber = protocol.PacketNumber(packetNumber)

	return header, nil
}

// GetLength gets the length of the publicHeader in bytes
// can only be called for regular packets
func (h *publicHeader) GetLength() (protocol.ByteCount, error) {
	if h.VersionFlag || h.ResetFlag {
		return 0, errGetLengthOnlyForRegularPackets
	}

	length := protocol.ByteCount(1) // 1 byte for public flags
	if h.PacketNumberLen != protocol.PacketNumberLen1 && h.PacketNumberLen != protocol.PacketNumberLen2 && h.PacketNumberLen != protocol.PacketNumberLen4 && h.PacketNumberLen != protocol.PacketNumberLen6 {
		return 0, errPacketNumberLenNotSet
	}
	if !h.TruncateConnectionID {
		length += 8 // 8 bytes for the connection ID
	}
	length += protocol.ByteCount(len(h.DiversificationNonce))
	length += protocol.ByteCount(h.PacketNumberLen)
	return length, nil
}
<|endoftext|>


<commit_before>package bootstrapCmd

import (
	// Stdlib
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"sort"

	// Internal
	"github.com/salsaflow/salsaflow/action"
	"github.com/salsaflow/salsaflow/app"
	"github.com/salsaflow/salsaflow/app/appflags"
	"github.com/salsaflow/salsaflow/config"
	"github.com/salsaflow/salsaflow/config/loader"
	"github.com/salsaflow/salsaflow/errs"
	"github.com/salsaflow/salsaflow/log"
	"github.com/salsaflow/salsaflow/modules"
	"github.com/salsaflow/salsaflow/prompt"

	// Other
	"gopkg.in/tchap/gocli.v2"
)

var Command = &gocli.Command{
	UsageLine: `
  bootstrap -skeleton=SKELETON [-skeleton_only]

  bootstrap -no_skeleton`,
	Short: "bootstrap repository for SalsaFlow",
	Long: `
  Bootstrap the repository for SalsaFlow.

  This command should be used to set up the local configuration directory
  for SalsaFlow (the directory that is then committed into the repository).

  The user is prompted for all necessary data.

  The -skeleton flag can be used to specify the repository to be used
  for custom scripts. It expects a string of "$OWNER/$REPO" and then uses
  the repository located at github.com/$OWNER/$REPO. It clones the repository
  and copies the content into the local configuration directory.

  In case no skeleton is to be used to bootstrap the repository,
  -no_skeleton must be specified explicitly.

  In case the repository is bootstrapped, but the skeleton is missing,
  it can be added by specifying -skeleton=SKELETON -skeleton_only.
  That will skip the configuration file generation step.
	`,
	Action: run,
}

var (
	flagNoSkeleton   bool
	flagSkeleton     string
	flagSkeletonOnly bool
)

func init() {
	// Register flags.
	Command.Flags.BoolVar(&flagNoSkeleton, "no_skeleton", flagNoSkeleton,
		"do not use any skeleton to bootstrap the repository")
	Command.Flags.StringVar(&flagSkeleton, "skeleton", flagSkeleton,
		"skeleton to be used to bootstrap the repository")
	Command.Flags.BoolVar(&flagSkeletonOnly, "skeleton_only", flagSkeletonOnly,
		"skip the config dialog and only install the skeleton")

	// Register global flags.
	appflags.RegisterGlobalFlags(&Command.Flags)
}

func run(cmd *gocli.Command, args []string) {
	if len(args) != 0 {
		cmd.Usage()
		os.Exit(2)
	}

	app.InitLogging()

	defer prompt.RecoverCancel()

	if err := runMain(cmd); err != nil {
		errs.Fatal(err)
	}
}

func runMain(cmd *gocli.Command) (err error) {
	// Validate CL flags.
	task := "Check the command line flags"
	switch {
	case flagSkeleton == "" && !flagNoSkeleton:
		cmd.Usage()
		return errs.NewError(
			task, errors.New("-no_skeleton must be specified when no skeleton is given"))

	case flagSkeletonOnly && flagSkeleton == "":
		cmd.Usage()
		return errs.NewError(
			task, errors.New("-skeleton must be specified when -skeleton_only is set"))
	}

	// Make sure the local config directory exists.
	if err := ensureLocalConfigDirectoryExists(); err != nil {
		return err
	}
	action.RollbackOnError(&err, action.ActionFunc(deleteLocalConfigDirectory))

	// Set up the global and local configuration file unless -skeleton_only.
	if !flagSkeletonOnly {
		if err := assembleAndWriteConfig(); err != nil {
			return err
		}
	}

	// Install the skeleton into the local config directory if desired.
	if skeleton := flagSkeleton; skeleton != "" {
		if err := getAndPourSkeleton(skeleton); err != nil {
			return err
		}
	}

	fmt.Println()
	log.Log("Successfully bootstrapped the repository for SalsaFlow")
	log.NewLine("Do not forget to commit modified configuration files!")
	return nil
}

func ensureLocalConfigDirectoryExists() error {
	task := "Make sure the local configuration directory exists"

	// Get the directory absolute path.
	localConfigDir, err := config.LocalConfigDirectoryAbsolutePath()
	if err != nil {
		return errs.NewError(task, err)
	}

	// In case the path exists, make sure it is a directory.
	info, err := os.Stat(localConfigDir)
	if err != nil {
		if !os.IsNotExist(err) {
			return errs.NewError(task, err)
		}
	} else {
		if !info.IsDir() {
			return errs.NewError(task, fmt.Errorf("not a directory: %v", localConfigDir))
		}
		return nil
	}

	// Otherwise create the directory.
	if err := os.MkdirAll(localConfigDir, 0755); err != nil {
		return errs.NewError(task, err)
	}

	return nil
}

func deleteLocalConfigDirectory() error {
	task := "Delete the local configuration directory"
	log.Run(task)

	// Get the directory absolute path.
	localConfigDir, err := config.LocalConfigDirectoryAbsolutePath()
	if err != nil {
		return errs.NewError(task, err)
	}

	// Delete the directory.
	if err := os.RemoveAll(localConfigDir); err != nil {
		return errs.NewError(task, err)
	}

	return nil
}

func assembleAndWriteConfig() error {
	// Group available modules by kind.
	var (
		issueTrackingModules []loader.Module
		codeReviewModules    []loader.Module
		releaseNotesModules  []loader.Module
	)
	groups := groupModulesByKind(modules.AvailableModules())
	for _, group := range groups {
		switch group[0].Kind() {
		case loader.ModuleKindIssueTracking:
			issueTrackingModules = group
		case loader.ModuleKindCodeReview:
			codeReviewModules = group
		case loader.ModuleKindReleaseNotes:
			releaseNotesModules = group
		}
	}

	sort.Sort(commonModules(issueTrackingModules))
	sort.Sort(commonModules(codeReviewModules))
	sort.Sort(commonModules(releaseNotesModules))

	// Run the common dialog.
	task := "Run the core configuration dialog"
	if err := loader.RunCommonBootstrapDialog(); err != nil {
		return errs.NewError(task, err)
	}

	// Run the dialog.
	task = "Run the modules configuration dialog"
	err := loader.RunModuleBootstrapDialog(
		&loader.ModuleDialogSection{issueTrackingModules, false},
		&loader.ModuleDialogSection{codeReviewModules, false},
		&loader.ModuleDialogSection{releaseNotesModules, true},
	)
	if err != nil {
		return errs.NewError(task, err)
	}

	return nil
}

func getAndPourSkeleton(skeleton string) error {
	// Get or update given skeleton.
	task := fmt.Sprintf("Get or update skeleton '%v'", skeleton)
	log.Run(task)
	if err := getOrUpdateSkeleton(flagSkeleton); err != nil {
		return errs.NewError(task, err)
	}

	// Move the skeleton files into place.
	task = "Copy the skeleton into the configuration directory"
	log.Go(task)

	localConfigDir, err := config.LocalConfigDirectoryAbsolutePath()
	if err != nil {
		return errs.NewError(task, err)
	}

	log.NewLine("")
	if err := pourSkeleton(flagSkeleton, localConfigDir); err != nil {
		return errs.NewError(task, err)
	}
	log.NewLine("")
	log.Ok(task)

	return nil
}

func writeConfigFile(path string, configObject interface{}) error {
	file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0640)
	if err != nil {
		return err
	}
	defer file.Close()

	content, err := config.Marshal(configObject)
	if err != nil {
		return err
	}

	if _, err := io.Copy(file, bytes.NewReader(content)); err != nil {
		return err
	}

	return nil
}
<commit_msg>repo bootstrap: Refactor<commit_after>package bootstrapCmd

import (
	// Stdlib
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"sort"

	// Internal
	"github.com/salsaflow/salsaflow/action"
	"github.com/salsaflow/salsaflow/app"
	"github.com/salsaflow/salsaflow/app/appflags"
	"github.com/salsaflow/salsaflow/config"
	"github.com/salsaflow/salsaflow/config/loader"
	"github.com/salsaflow/salsaflow/errs"
	"github.com/salsaflow/salsaflow/log"
	"github.com/salsaflow/salsaflow/modules"
	"github.com/salsaflow/salsaflow/prompt"

	// Other
	"gopkg.in/tchap/gocli.v2"
)

var Command = &gocli.Command{
	UsageLine: `
  bootstrap -skeleton=SKELETON [-skeleton_only]

  bootstrap -no_skeleton`,
	Short: "bootstrap repository for SalsaFlow",
	Long: `
  Bootstrap the repository for SalsaFlow.

  This command should be used to set up the local configuration directory
  for SalsaFlow (the directory that is then committed into the repository).

  The user is prompted for all necessary data.

  The -skeleton flag can be used to specify the repository to be used
  for custom scripts. It expects a string of "$OWNER/$REPO" and then uses
  the repository located at github.com/$OWNER/$REPO. It clones the repository
  and copies the content into the local configuration directory.

  In case no skeleton is to be used to bootstrap the repository,
  -no_skeleton must be specified explicitly.

  In case the repository is bootstrapped, but the skeleton is missing,
  it can be added by specifying -skeleton=SKELETON -skeleton_only.
  That will skip the configuration file generation step.
	`,
	Action: run,
}

var (
	flagNoSkeleton   bool
	flagSkeleton     string
	flagSkeletonOnly bool
)

func init() {
	// Register flags.
	Command.Flags.BoolVar(&flagNoSkeleton, "no_skeleton", flagNoSkeleton,
		"do not use any skeleton to bootstrap the repository")
	Command.Flags.StringVar(&flagSkeleton, "skeleton", flagSkeleton,
		"skeleton to be used to bootstrap the repository")
	Command.Flags.BoolVar(&flagSkeletonOnly, "skeleton_only", flagSkeletonOnly,
		"skip the config dialog and only install the skeleton")

	// Register global flags.
	appflags.RegisterGlobalFlags(&Command.Flags)
}

func run(cmd *gocli.Command, args []string) {
	if len(args) != 0 {
		cmd.Usage()
		os.Exit(2)
	}

	app.InitLogging()

	defer prompt.RecoverCancel()

	if err := runMain(cmd); err != nil {
		errs.Fatal(err)
	}
}

func runMain(cmd *gocli.Command) (err error) {
	// Validate CL flags.
	task := "Check the command line flags"
	switch {
	case flagSkeleton == "" && !flagNoSkeleton:
		cmd.Usage()
		return errs.NewError(
			task, errors.New("-no_skeleton must be specified when no skeleton is given"))

	case flagSkeletonOnly && flagSkeleton == "":
		cmd.Usage()
		return errs.NewError(
			task, errors.New("-skeleton must be specified when -skeleton_only is set"))
	}

	// Make sure the local config directory exists.
	act, err := ensureLocalConfigDirectoryExists()
	if err != nil {
		return err
	}
	action.RollbackOnError(&err, act)

	// Set up the global and local configuration file unless -skeleton_only.
	if !flagSkeletonOnly {
		if err := assembleAndWriteConfig(); err != nil {
			return err
		}
	}

	// Install the skeleton into the local config directory if desired.
	if skeleton := flagSkeleton; skeleton != "" {
		if err := getAndPourSkeleton(skeleton); err != nil {
			return err
		}
	}

	fmt.Println()
	log.Log("Successfully bootstrapped the repository for SalsaFlow")
	log.NewLine("Do not forget to commit modified configuration files!")
	return nil
}

func ensureLocalConfigDirectoryExists() (action.Action, error) {
	task := "Make sure the local configuration directory exists"

	// Get the directory absolute path.
	localConfigDir, err := config.LocalConfigDirectoryAbsolutePath()
	if err != nil {
		return nil, errs.NewError(task, err)
	}

	// In case the path exists, make sure it is a directory.
	info, err := os.Stat(localConfigDir)
	if err != nil {
		if !os.IsNotExist(err) {
			return nil, errs.NewError(task, err)
		}
	} else {
		if !info.IsDir() {
			return nil, errs.NewError(task, fmt.Errorf("not a directory: %v", localConfigDir))
		}
		return action.Noop, nil
	}

	// Otherwise create the directory.
	if err := os.MkdirAll(localConfigDir, 0755); err != nil {
		return nil, errs.NewError(task, err)
	}

	// Return the rollback function.
	act := action.ActionFunc(func() error {
		// Delete the directory.
		log.Rollback(task)
		task := "Delete the local configuration directory"
		if err := os.RemoveAll(localConfigDir); err != nil {
			return errs.NewError(task, err)
		}
		return nil
	})
	return act, nil
}

func assembleAndWriteConfig() error {
	// Group available modules by kind.
	var (
		issueTrackingModules []loader.Module
		codeReviewModules    []loader.Module
		releaseNotesModules  []loader.Module
	)
	groups := groupModulesByKind(modules.AvailableModules())
	for _, group := range groups {
		switch group[0].Kind() {
		case loader.ModuleKindIssueTracking:
			issueTrackingModules = group
		case loader.ModuleKindCodeReview:
			codeReviewModules = group
		case loader.ModuleKindReleaseNotes:
			releaseNotesModules = group
		}
	}

	sort.Sort(commonModules(issueTrackingModules))
	sort.Sort(commonModules(codeReviewModules))
	sort.Sort(commonModules(releaseNotesModules))

	// Run the common dialog.
	task := "Run the core configuration dialog"
	if err := loader.RunCommonBootstrapDialog(); err != nil {
		return errs.NewError(task, err)
	}

	// Run the dialog.
	task = "Run the modules configuration dialog"
	err := loader.RunModuleBootstrapDialog(
		&loader.ModuleDialogSection{issueTrackingModules, false},
		&loader.ModuleDialogSection{codeReviewModules, false},
		&loader.ModuleDialogSection{releaseNotesModules, true},
	)
	if err != nil {
		return errs.NewError(task, err)
	}

	return nil
}

func getAndPourSkeleton(skeleton string) error {
	// Get or update given skeleton.
	task := fmt.Sprintf("Get or update skeleton '%v'", skeleton)
	log.Run(task)
	if err := getOrUpdateSkeleton(flagSkeleton); err != nil {
		return errs.NewError(task, err)
	}

	// Move the skeleton files into place.
	task = "Copy the skeleton into the configuration directory"
	log.Go(task)

	localConfigDir, err := config.LocalConfigDirectoryAbsolutePath()
	if err != nil {
		return errs.NewError(task, err)
	}

	log.NewLine("")
	if err := pourSkeleton(flagSkeleton, localConfigDir); err != nil {
		return errs.NewError(task, err)
	}
	log.NewLine("")
	log.Ok(task)

	return nil
}

func writeConfigFile(path string, configObject interface{}) error {
	file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0640)
	if err != nil {
		return err
	}
	defer file.Close()

	content, err := config.Marshal(configObject)
	if err != nil {
		return err
	}

	if _, err := io.Copy(file, bytes.NewReader(content)); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package socks

import (
	"io"
	"net"
	"strconv"
	"sync"
	"time"

	"github.com/v2ray/v2ray-core"
	"github.com/v2ray/v2ray-core/common/alloc"
	"github.com/v2ray/v2ray-core/common/errors"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	jsonconfig "github.com/v2ray/v2ray-core/proxy/socks/config/json"
	"github.com/v2ray/v2ray-core/proxy/socks/protocol"
)

// SocksServer is a SOCKS 5 proxy server
type SocksServer struct {
	accepting bool
	vPoint    *core.Point
	config    *jsonconfig.SocksConfig
}

func NewSocksServer(vp *core.Point, config *jsonconfig.SocksConfig) *SocksServer {
	return &SocksServer{
		vPoint: vp,
		config: config,
	}
}

func (server *SocksServer) Listen(port uint16) error {
	listener, err := net.Listen("tcp", ":"+strconv.Itoa(int(port)))
	if err != nil {
		log.Error("Socks failed to listen on port %d: %v", port, err)
		return err
	}
	server.accepting = true
	go server.AcceptConnections(listener)
	if server.config.UDPEnabled {
		server.ListenUDP(port)
	}
	return nil
}

func (server *SocksServer) AcceptConnections(listener net.Listener) {
	for server.accepting {
		connection, err := listener.Accept()
		if err != nil {
			log.Error("Socks failed to accept new connection %v", err)
			return
		}
		go server.HandleConnection(connection)
	}
}

func (server *SocksServer) HandleConnection(connection net.Conn) error {
	defer connection.Close()

	reader := v2net.NewTimeOutReader(120, connection)

	auth, auth4, err := protocol.ReadAuthentication(reader)
	if err != nil && !errors.HasCode(err, 1000) {
		log.Error("Socks failed to read authentication: %v", err)
		return err
	}

	if err != nil && errors.HasCode(err, 1000) {
		return server.handleSocks4(reader, connection, auth4)
	} else {
		return server.handleSocks5(reader, connection, auth)
	}
}

func (server *SocksServer) handleSocks5(reader *v2net.TimeOutReader, writer io.Writer, auth protocol.Socks5AuthenticationRequest) error {
	expectedAuthMethod := protocol.AuthNotRequired
	if server.config.IsPassword() {
		expectedAuthMethod = protocol.AuthUserPass
	}

	if !auth.HasAuthMethod(expectedAuthMethod) {
		authResponse := protocol.NewAuthenticationResponse(protocol.AuthNoMatchingMethod)
		err := protocol.WriteAuthentication(writer, authResponse)
		if err != nil {
			log.Error("Socks failed to write authentication: %v", err)
			return err
		}
		log.Warning("Socks client doesn't support allowed any auth methods.")
		return errors.NewInvalidOperationError("Unsupported auth methods.")
	}

	authResponse := protocol.NewAuthenticationResponse(expectedAuthMethod)
	err := protocol.WriteAuthentication(writer, authResponse)
	if err != nil {
		log.Error("Socks failed to write authentication: %v", err)
		return err
	}
	if server.config.IsPassword() {
		upRequest, err := protocol.ReadUserPassRequest(reader)
		if err != nil {
			log.Error("Socks failed to read username and password: %v", err)
			return err
		}
		status := byte(0)
		if server.config.HasAccount(upRequest.Username(), upRequest.Password()) {
			status = byte(0xFF)
		}
		upResponse := protocol.NewSocks5UserPassResponse(status)
		err = protocol.WriteUserPassResponse(writer, upResponse)
		if err != nil {
			log.Error("Socks failed to write user pass response: %v", err)
			return err
		}
		if status != byte(0) {
			err = errors.NewAuthenticationError(upRequest.AuthDetail())
			log.Warning(err.Error())
			return err
		}
	}

	request, err := protocol.ReadRequest(reader)
	if err != nil {
		log.Error("Socks failed to read request: %v", err)
		return err
	}

	if request.Command == protocol.CmdUdpAssociate && server.config.UDPEnabled {
		return server.handleUDP(reader, writer)
	}

	response := protocol.NewSocks5Response()
	if request.Command == protocol.CmdBind || request.Command == protocol.CmdUdpAssociate {
		response := protocol.NewSocks5Response()
		response.Error = protocol.ErrorCommandNotSupported

		responseBuffer := alloc.NewSmallBuffer().Clear()
		response.Write(responseBuffer)
		_, err = writer.Write(responseBuffer.Value)
		if err != nil {
			log.Error("Socks failed to write response: %v", err)
			return err
		}
		log.Warning("Unsupported socks command %d", request.Command)
		return errors.NewInvalidOperationError("Socks command " + strconv.Itoa(int(request.Command)))
	}

	response.Error = protocol.ErrorSuccess

	// Some SOCKS software requires a value other than dest. Let's fake one:
	response.Port = uint16(1717)
	response.AddrType = protocol.AddrTypeIPv4
	response.IPv4[0] = 0
	response.IPv4[1] = 0
	response.IPv4[2] = 0
	response.IPv4[3] = 0

	responseBuffer := alloc.NewSmallBuffer().Clear()
	response.Write(responseBuffer)
	_, err = writer.Write(responseBuffer.Value)
	if err != nil {
		log.Error("Socks failed to write response: %v", err)
		return err
	}

	dest := request.Destination()
	data, err := v2net.ReadFrom(reader, nil)
	if err != nil {
		return err
	}

	packet := v2net.NewPacket(dest, data, true)
	server.transport(reader, writer, packet)
	return nil
}

func (server *SocksServer) handleUDP(reader *v2net.TimeOutReader, writer io.Writer) error {
	response := protocol.NewSocks5Response()
	response.Error = protocol.ErrorSuccess

	udpAddr := server.getUDPAddr()

	response.Port = udpAddr.Port()
	switch {
	case udpAddr.IsIPv4():
		response.AddrType = protocol.AddrTypeIPv4
		copy(response.IPv4[:], udpAddr.IP())
	case udpAddr.IsIPv6():
		response.AddrType = protocol.AddrTypeIPv6
		copy(response.IPv6[:], udpAddr.IP())
	case udpAddr.IsDomain():
		response.AddrType = protocol.AddrTypeDomain
		response.Domain = udpAddr.Domain()
	}

	responseBuffer := alloc.NewSmallBuffer()
	response.Write(responseBuffer)
	_, err := writer.Write(responseBuffer.Value)
	responseBuffer.Release()

	if err != nil {
		log.Error("Socks failed to write response: %v", err)
		return err
	}

	reader.SetTimeOut(300)      /* 5 minutes */
	v2net.ReadFrom(reader, nil) // Just in case of anything left in the socket
	// The TCP connection closes after this method returns. We need to wait until
	// the client closes it.
	// TODO: get notified from UDP part
	<-time.After(5 * time.Minute)

	return nil
}

func (server *SocksServer) handleSocks4(reader io.Reader, writer io.Writer, auth protocol.Socks4AuthenticationRequest) error {
	result := protocol.Socks4RequestGranted
	if auth.Command == protocol.CmdBind {
		result = protocol.Socks4RequestRejected
	}
	socks4Response := protocol.NewSocks4AuthenticationResponse(result, auth.Port, auth.IP[:])

	responseBuffer := alloc.NewSmallBuffer().Clear()
	socks4Response.Write(responseBuffer)
	writer.Write(responseBuffer.Value)
	responseBuffer.Release()

	if result == protocol.Socks4RequestRejected {
		return errors.NewInvalidOperationError("Socks4 command " + strconv.Itoa(int(auth.Command)))
	}

	dest := v2net.NewTCPDestination(v2net.IPAddress(auth.IP[:], auth.Port))
	data, err := v2net.ReadFrom(reader, nil)
	if err != nil {
		return err
	}

	packet := v2net.NewPacket(dest, data, true)
	server.transport(reader, writer, packet)
	return nil
}

func (server *SocksServer) transport(reader io.Reader, writer io.Writer, firstPacket v2net.Packet) {
	ray := server.vPoint.DispatchToOutbound(firstPacket)
	input := ray.InboundInput()
	output := ray.InboundOutput()

	var inputFinish, outputFinish sync.Mutex
	inputFinish.Lock()
	outputFinish.Lock()

	go dumpInput(reader, input, &inputFinish)
	go dumpOutput(writer, output, &outputFinish)
	outputFinish.Lock()
}

func dumpInput(reader io.Reader, input chan<- *alloc.Buffer, finish *sync.Mutex) {
	v2net.ReaderToChan(input, reader)
	finish.Unlock()
	close(input)
}

func dumpOutput(writer io.Writer, output <-chan *alloc.Buffer, finish *sync.Mutex) {
	v2net.ChanToWriter(writer, output)
	finish.Unlock()
}
<commit_msg>release buffer after using<commit_after>package socks

import (
	"io"
	"net"
	"strconv"
	"sync"
	"time"

	"github.com/v2ray/v2ray-core"
	"github.com/v2ray/v2ray-core/common/alloc"
	"github.com/v2ray/v2ray-core/common/errors"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	jsonconfig "github.com/v2ray/v2ray-core/proxy/socks/config/json"
	"github.com/v2ray/v2ray-core/proxy/socks/protocol"
)

// SocksServer is a SOCKS 5 proxy server
type SocksServer struct {
	accepting bool
	vPoint    *core.Point
	config    *jsonconfig.SocksConfig
}

func NewSocksServer(vp *core.Point, config *jsonconfig.SocksConfig) *SocksServer {
	return &SocksServer{
		vPoint: vp,
		config: config,
	}
}

func (server *SocksServer) Listen(port uint16) error {
	listener, err := net.Listen("tcp", ":"+strconv.Itoa(int(port)))
	if err != nil {
		log.Error("Socks failed to listen on port %d: %v", port, err)
		return err
	}
	server.accepting = true
	go server.AcceptConnections(listener)
	if server.config.UDPEnabled {
		server.ListenUDP(port)
	}
	return nil
}

func (server *SocksServer) AcceptConnections(listener net.Listener) {
	for server.accepting {
		connection, err := listener.Accept()
		if err != nil {
			log.Error("Socks failed to accept new connection %v", err)
			return
		}
		go server.HandleConnection(connection)
	}
}

func (server *SocksServer) HandleConnection(connection net.Conn) error {
	defer connection.Close()

	reader := v2net.NewTimeOutReader(120, connection)

	auth, auth4, err := protocol.ReadAuthentication(reader)
	if err != nil && !errors.HasCode(err, 1000) {
		log.Error("Socks failed to read authentication: %v", err)
		return err
	}

	if err != nil && errors.HasCode(err, 1000) {
		return server.handleSocks4(reader, connection, auth4)
	} else {
		return server.handleSocks5(reader, connection, auth)
	}
}

func (server *SocksServer) handleSocks5(reader *v2net.TimeOutReader, writer io.Writer, auth protocol.Socks5AuthenticationRequest) error {
	expectedAuthMethod := protocol.AuthNotRequired
	if server.config.IsPassword() {
		expectedAuthMethod = protocol.AuthUserPass
	}

	if !auth.HasAuthMethod(expectedAuthMethod) {
		authResponse := protocol.NewAuthenticationResponse(protocol.AuthNoMatchingMethod)
		err := protocol.WriteAuthentication(writer, authResponse)
		if err != nil {
			log.Error("Socks failed to write authentication: %v", err)
			return err
		}
		log.Warning("Socks client doesn't support allowed any auth methods.")
		return errors.NewInvalidOperationError("Unsupported auth methods.")
	}

	authResponse := protocol.NewAuthenticationResponse(expectedAuthMethod)
	err := protocol.WriteAuthentication(writer, authResponse)
	if err != nil {
		log.Error("Socks failed to write authentication: %v", err)
		return err
	}
	if server.config.IsPassword() {
		upRequest, err := protocol.ReadUserPassRequest(reader)
		if err != nil {
			log.Error("Socks failed to read username and password: %v", err)
			return err
		}
		status := byte(0)
		if server.config.HasAccount(upRequest.Username(), upRequest.Password()) {
			status = byte(0xFF)
		}
		upResponse := protocol.NewSocks5UserPassResponse(status)
		err = protocol.WriteUserPassResponse(writer, upResponse)
		if err != nil {
			log.Error("Socks failed to write user pass response: %v", err)
			return err
		}
		if status != byte(0) {
			err = errors.NewAuthenticationError(upRequest.AuthDetail())
			log.Warning(err.Error())
			return err
		}
	}

	request, err := protocol.ReadRequest(reader)
	if err != nil {
		log.Error("Socks failed to read request: %v", err)
		return err
	}

	if request.Command == protocol.CmdUdpAssociate && server.config.UDPEnabled {
		return server.handleUDP(reader, writer)
	}

	response := protocol.NewSocks5Response()
	if request.Command == protocol.CmdBind || request.Command == protocol.CmdUdpAssociate {
		response := protocol.NewSocks5Response()
		response.Error = protocol.ErrorCommandNotSupported

		responseBuffer := alloc.NewSmallBuffer().Clear()
		response.Write(responseBuffer)
		_, err = writer.Write(responseBuffer.Value)
		responseBuffer.Release()
		if err != nil {
			log.Error("Socks failed to write response: %v", err)
			return err
		}
		log.Warning("Unsupported socks command %d", request.Command)
		return errors.NewInvalidOperationError("Socks command " + strconv.Itoa(int(request.Command)))
	}

	response.Error = protocol.ErrorSuccess

	// Some SOCKS software requires a value other than dest. Let's fake one:
	response.Port = uint16(1717)
	response.AddrType = protocol.AddrTypeIPv4
	response.IPv4[0] = 0
	response.IPv4[1] = 0
	response.IPv4[2] = 0
	response.IPv4[3] = 0

	responseBuffer := alloc.NewSmallBuffer().Clear()
	response.Write(responseBuffer)
	_, err = writer.Write(responseBuffer.Value)
	responseBuffer.Release()
	if err != nil {
		log.Error("Socks failed to write response: %v", err)
		return err
	}

	dest := request.Destination()
	data, err := v2net.ReadFrom(reader, nil)
	if err != nil {
		return err
	}

	packet := v2net.NewPacket(dest, data, true)
	server.transport(reader, writer, packet)
	return nil
}

func (server *SocksServer) handleUDP(reader *v2net.TimeOutReader, writer io.Writer) error {
	response := protocol.NewSocks5Response()
	response.Error = protocol.ErrorSuccess

	udpAddr := server.getUDPAddr()

	response.Port = udpAddr.Port()
	switch {
	case udpAddr.IsIPv4():
		response.AddrType = protocol.AddrTypeIPv4
		copy(response.IPv4[:], udpAddr.IP())
	case udpAddr.IsIPv6():
		response.AddrType = protocol.AddrTypeIPv6
		copy(response.IPv6[:], udpAddr.IP())
	case udpAddr.IsDomain():
		response.AddrType = protocol.AddrTypeDomain
		response.Domain = udpAddr.Domain()
	}

	responseBuffer := alloc.NewSmallBuffer()
	response.Write(responseBuffer)
	_, err := writer.Write(responseBuffer.Value)
	responseBuffer.Release()

	if err != nil {
		log.Error("Socks failed to write response: %v", err)
		return err
	}

	reader.SetTimeOut(300)      /* 5 minutes */
	v2net.ReadFrom(reader, nil) // Just in case of anything left in the socket
	// The TCP connection closes after this method returns. We need to wait until
	// the client closes it.
	// TODO: get notified from UDP part
	<-time.After(5 * time.Minute)

	return nil
}

func (server *SocksServer) handleSocks4(reader io.Reader, writer io.Writer, auth protocol.Socks4AuthenticationRequest) error {
	result := protocol.Socks4RequestGranted
	if auth.Command == protocol.CmdBind {
		result = protocol.Socks4RequestRejected
	}
	socks4Response := protocol.NewSocks4AuthenticationResponse(result, auth.Port, auth.IP[:])

	responseBuffer := alloc.NewSmallBuffer().Clear()
	socks4Response.Write(responseBuffer)
	writer.Write(responseBuffer.Value)
	responseBuffer.Release()

	if result == protocol.Socks4RequestRejected {
		return errors.NewInvalidOperationError("Socks4 command " + strconv.Itoa(int(auth.Command)))
	}

	dest := v2net.NewTCPDestination(v2net.IPAddress(auth.IP[:], auth.Port))
	data, err := v2net.ReadFrom(reader, nil)
	if err != nil {
		return err
	}

	packet := v2net.NewPacket(dest, data, true)
	server.transport(reader, writer, packet)
	return nil
}

func (server *SocksServer) transport(reader io.Reader, writer io.Writer, firstPacket v2net.Packet) {
	ray := server.vPoint.DispatchToOutbound(firstPacket)
	input := ray.InboundInput()
	output := ray.InboundOutput()

	var inputFinish, outputFinish sync.Mutex
	inputFinish.Lock()
	outputFinish.Lock()

	go dumpInput(reader, input, &inputFinish)
	go dumpOutput(writer, output, &outputFinish)
	outputFinish.Lock()
}

func dumpInput(reader io.Reader, input chan<- *alloc.Buffer, finish *sync.Mutex) {
	v2net.ReaderToChan(input, reader)
	finish.Unlock()
	close(input)
}

func dumpOutput(writer io.Writer, output <-chan *alloc.Buffer, finish *sync.Mutex) {
	v2net.ChanToWriter(writer, output)
	finish.Unlock()
}
<|endoftext|>