<commit_before>package ydls

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sort"
	"sync"

	"github.com/wader/ydls/ffmpeg"
	"github.com/wader/ydls/youtubedl"
)

func writeID3v2FromYoutueDLInfo(w io.Writer, i *youtubedl.Info) {
	id3v2Frames := []id3v2Frame{
		&textFrame{"TPE1", firstNonEmpty(i.Artist, i.Creator, i.Uploader)},
		&textFrame{"TIT2", i.Title},
		&commFrame{"XXX", "", i.Description},
	}
	if i.Duration > 0 {
		id3v2Frames = append(id3v2Frames, &textFrame{"TLEN", fmt.Sprintf("%d", uint32(i.Duration*1000))})
	}
	if len(i.ThumbnailBytes) > 0 {
		id3v2Frames = append(id3v2Frames, &apicFrame{
			http.DetectContentType(i.ThumbnailBytes),
			id3v2PictureTypeOther,
			"",
			i.ThumbnailBytes,
		})
	}

	id3v2WriteHeader(w, id3v2Frames)
}

func findFormat(formats []*youtubedl.Format, protocol string, aCodecs *prioStringSet, vCodecs *prioStringSet) *youtubedl.Format {
	var matched []*youtubedl.Format

	for _, f := range formats {
		if protocol != "*" && f.Protocol != protocol {
			continue
		}
		if !(aCodecs == nil || (f.NormACodec == "" && aCodecs.empty()) || aCodecs.member(f.NormACodec)) {
			continue
		}
		if !(vCodecs == nil || (f.NormVCodec == "" && vCodecs.empty()) || vCodecs.member(f.NormVCodec)) {
			continue
		}

		matched = append(matched, f)
	}

	sort.Sort(youtubedl.FormatByNormBR(matched))

	if len(matched) > 0 {
		return matched[0]
	}

	return nil
}

func findBestFormats(ydlFormats []*youtubedl.Format, format *Format) (aFormat *youtubedl.Format, vFormat *youtubedl.Format) {
	type neededFormat struct {
		aCodecs    *prioStringSet
		vCodecs    *prioStringSet
		aYDLFormat **youtubedl.Format
		vYDLFormat **youtubedl.Format
	}

	var neededFormats []*neededFormat

	// match exactly, both audio/video codecs found or not found
	neededFormats = append(neededFormats, &neededFormat{&format.ACodecs, &format.VCodecs, &aFormat, &vFormat})

	if !format.ACodecs.empty() {
		// matching audio codec with any video codec
		neededFormats = append(neededFormats, &neededFormat{&format.ACodecs, nil, &aFormat, nil})
		// match any audio codec and no video
		neededFormats = append(neededFormats, &neededFormat{nil, &prioStringSet{}, &aFormat, nil})
		// match any audio and video codec
		neededFormats = append(neededFormats, &neededFormat{nil, nil, &aFormat, nil})
	}
	if !format.VCodecs.empty() {
		// same logic as above
		neededFormats = append(neededFormats, &neededFormat{nil, &format.VCodecs, nil, &vFormat})
		neededFormats = append(neededFormats, &neededFormat{&prioStringSet{}, nil, nil, &vFormat})
		neededFormats = append(neededFormats, &neededFormat{nil, nil, nil, &vFormat})
	}

	// TODO: if only audio... choose stream with lowest video br?
	// TODO: rtmp only if mp4 etc?

	// prefer rtmp as it seems to send fragmented mp4 etc
	for _, proto := range []string{"rtmp", "*"} {
		for _, f := range neededFormats {
			m := findFormat(ydlFormats, proto, f.aCodecs, f.vCodecs)

			if m == nil {
				continue
			}

			if f.aYDLFormat != nil && *f.aYDLFormat == nil && m.NormACodec != "" {
				*f.aYDLFormat = m
			}
			if f.vYDLFormat != nil && *f.vYDLFormat == nil && m.NormVCodec != "" {
				*f.vYDLFormat = m
			}

			if (format.ACodecs.empty() || aFormat != nil) &&
				(format.VCodecs.empty() || vFormat != nil) {
				break
			}
		}
	}

	return aFormat, vFormat
}

func downloadAndProbeFormat(ydl *youtubedl.Info, filter string, debugLog *log.Logger) (r io.ReadCloser, pi *ffmpeg.ProbeInfo, err error) {
	var ydlStderr io.Writer
	if debugLog != nil {
		ydlStderr = &loggerWriter{Logger: debugLog, Prefix: "ydl 2> "}
	}
	r, err = ydl.Download(filter, ydlStderr)
	if err != nil {
		return nil, nil, err
	}

	rr := &reReadCloser{ReadCloser: r}

	var ffprobeStderr io.Writer
	if debugLog != nil {
		ffprobeStderr = &loggerWriter{Logger: debugLog, Prefix: fmt.Sprintf("ffprobe %s 2> ", filter)}
	}
	const maxProbeByteSize = 10 * 1024 * 1024
	pi, err = ffmpeg.FFprobe(io.LimitReader(rr, maxProbeByteSize), debugLog, ffprobeStderr)
	if err != nil {
		return nil, nil, err
	}
	// restart and replay buffer data used when probing
	rr.Restarted = true

	return rr, pi, nil
}

// YDLs youtubedl downloader with some extras
type YDLs struct {
	Formats Formats
}

// NewFromFile new YDLs using formats file
func NewFromFile(formatsPath string) (*YDLs, error) {
	formatsFile, err := os.Open(formatsPath)
	if err != nil {
		return nil, err
	}
	defer formatsFile.Close()
	formats, err := parseFormats(formatsFile)
	if err != nil {
		return nil, err
	}

	return &YDLs{Formats: formats}, nil
}

// Download downloads media from URL and makes sure output is in specified format
func (ydls *YDLs) Download(url string, formatName string, debugLog *log.Logger) (r io.ReadCloser, filename string, mimeType string, err error) {
	log := log.New(ioutil.Discard, "", 0)
	if debugLog != nil {
		log = debugLog
	}

	log.Printf("URL: %s", url)

	var ydlStdout io.Writer
	if debugLog != nil {
		ydlStdout = &loggerWriter{Logger: debugLog, Prefix: "ydl 1> "}
	}
	ydl, err := youtubedl.NewFromURL(url, ydlStdout)
	if err != nil {
		log.Printf("Failed to download: %s", err)
		return nil, "", "", fmt.Errorf("failed to download")
	}

	log.Printf("Title: %s", ydl.Title)
	log.Printf("Available youtubedl formats:")
	for _, f := range ydl.Formats {
		log.Printf("  %s", f)
	}

	var outFormat *Format

	if formatName == "" {
		var probedInfo *ffmpeg.ProbeInfo
		r, probedInfo, err = downloadAndProbeFormat(ydl, "best[protocol=rtmp]/best", debugLog)
		if err != nil {
			return nil, "", "", err
		}

		log.Printf("Probed format %s", probedInfo)

		// see if we know about the probed format, otherwise fallback to "raw"
		outFormat = ydls.Formats.Find(probedInfo.FormatName(), probedInfo.ACodec(), probedInfo.VCodec())
	} else {
		outFormat = ydls.Formats.FindByName(formatName)
		if outFormat == nil {
			return nil, "", "", fmt.Errorf("could not find format")
		}

		aYDLFormat, vYDLFormat := findBestFormats(ydl.Formats, outFormat)

		log.Printf("Best youtubedl match for %s a=%s v=%s", formatName, aYDLFormat, vYDLFormat)

		var aProbedFormat *ffmpeg.ProbeInfo
		var aReader io.ReadCloser
		var vProbedFormat *ffmpeg.ProbeInfo
		var vReader io.ReadCloser

		// FIXME: messy

		if aYDLFormat != nil && vYDLFormat != nil {
			if aYDLFormat != vYDLFormat {
				var probeWG sync.WaitGroup
				probeWG.Add(1)
				go func() {
					defer probeWG.Done()
					aReader, aProbedFormat, err = downloadAndProbeFormat(ydl, aYDLFormat.FormatID, debugLog)
				}()
				probeWG.Add(1)
				go func() {
					defer probeWG.Done()
					vReader, vProbedFormat, err = downloadAndProbeFormat(ydl, vYDLFormat.FormatID, debugLog)
				}()
				probeWG.Wait()
			} else {
				aReader, aProbedFormat, err = downloadAndProbeFormat(ydl, aYDLFormat.FormatID, debugLog)
				vReader, vProbedFormat = aReader, aProbedFormat
			}
		} else if aYDLFormat != nil && vYDLFormat == nil {
			aReader, aProbedFormat, err = downloadAndProbeFormat(ydl, aYDLFormat.FormatID, debugLog)
		} else {
			aReader, aProbedFormat, err = downloadAndProbeFormat(ydl, "best", debugLog)
			vReader, vProbedFormat = aReader, aProbedFormat
		}

		if err != nil {
			return nil, "", "", fmt.Errorf("failed to probe")
		}

		log.Printf("Stream mapping:")
		var maps []ffmpeg.Map
		if len(outFormat.ACodecs) > 0 && aProbedFormat != nil && aProbedFormat.ACodec() != "" {
			canCopy := outFormat.ACodecs.member(aProbedFormat.ACodec())
			ffmpegCodec := boolString(canCopy, "copy", outFormat.ACodecs.first())
			ydlACodec := "n/a"
			if aYDLFormat != nil {
				ydlACodec = aYDLFormat.NormACodec
			}

			maps = append(maps, ffmpeg.Map{
				Input:           aReader,
				Kind:            "audio",
				StreamSpecifier: "a:0",
				Codec:           ffmpegCodec,
				Flags:           outFormat.ACodecFlags,
			})
			log.Printf("  audio probed:%s ydl:%s -> %s", aProbedFormat.ACodec(), ydlACodec, ffmpegCodec)
		}
		if len(outFormat.VCodecs) > 0 && vProbedFormat != nil && vProbedFormat.VCodec() != "" {
			canCopy := outFormat.VCodecs.member(vProbedFormat.VCodec())
			ffmpegCodec := boolString(canCopy, "copy", outFormat.VCodecs.first())
			ydlVCodec := "n/a"
			if vYDLFormat != nil {
				ydlVCodec = vYDLFormat.NormVCodec
			}

			maps = append(maps, ffmpeg.Map{
				Input:           vReader,
				Kind:            "video",
				StreamSpecifier: "v:0",
				Codec:           ffmpegCodec,
				Flags:           outFormat.VCodecFlags,
			})
			log.Printf("  video probed:%s ydl:%s -> %s", vProbedFormat.VCodec(), ydlVCodec, ffmpegCodec)
		}

		var ffmpegStderr io.Writer
		if debugLog != nil {
			ffmpegStderr = &loggerWriter{Logger: debugLog, Prefix: "ffmpeg 2> "}
		}
		ffmpegR, ffmpegW := io.Pipe()

		f := &ffmpeg.FFmpeg{
			Maps:     maps,
			Format:   ffmpeg.Format{Name: outFormat.Formats.first(), Flags: outFormat.FormatFlags},
			DebugLog: debugLog,
			Stdout:   ffmpegW,
			Stderr:   ffmpegStderr,
		}

		closeReaders := func() {
			ffmpegR.Close()
			if aReader != nil {
				aReader.Close()
			}
			if vReader != nil {
				vReader.Close()
			}
		}

		if err := f.Start(); err != nil {
			return nil, "", "", err
		}

		// probe read one byte to see if ffmpeg is happy
		b := make([]byte, 1)
		if _, err := io.ReadFull(ffmpegR, b); err != nil {
			defer closeReaders()
			if err := f.Wait(); err != nil {
				log.Printf("ffmpeg failed: %s", err)
				return nil, "", "", err
			}
			log.Printf("read failed: %s", err)
			return nil, "", "", err
		}

		var w io.WriteCloser
		r, w = io.Pipe()
		go func() {
			defer closeReaders()
			defer w.Close()

			if outFormat.Prepend == "id3v2" {
				writeID3v2FromYoutueDLInfo(w, ydl)
			}
			w.Write(b)
			io.Copy(w, ffmpegR)
			f.Wait()
		}()
	}

	if outFormat == nil {
		mimeType = "application/octet-stream"
		filename = ydl.Title + ".raw"
	} else {
		mimeType = outFormat.MIMEType
		filename = ydl.Title + "." + outFormat.Ext
	}

	return r, filename, mimeType, nil
}
<commit_msg>Make sure to close readers and fix append race<commit_after>package ydls

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sort"
	"sync"

	"github.com/wader/ydls/ffmpeg"
	"github.com/wader/ydls/youtubedl"
)

func writeID3v2FromYoutueDLInfo(w io.Writer, i *youtubedl.Info) {
	id3v2Frames := []id3v2Frame{
		&textFrame{"TPE1", firstNonEmpty(i.Artist, i.Creator, i.Uploader)},
		&textFrame{"TIT2", i.Title},
		&commFrame{"XXX", "", i.Description},
	}
	if i.Duration > 0 {
		id3v2Frames = append(id3v2Frames, &textFrame{"TLEN", fmt.Sprintf("%d", uint32(i.Duration*1000))})
	}
	if len(i.ThumbnailBytes) > 0 {
		id3v2Frames = append(id3v2Frames, &apicFrame{
			http.DetectContentType(i.ThumbnailBytes),
			id3v2PictureTypeOther,
			"",
			i.ThumbnailBytes,
		})
	}

	id3v2WriteHeader(w, id3v2Frames)
}

func findFormat(formats []*youtubedl.Format, protocol string, aCodecs *prioStringSet, vCodecs *prioStringSet) *youtubedl.Format {
	var matched []*youtubedl.Format

	for _, f := range formats {
		if protocol != "*" && f.Protocol != protocol {
			continue
		}
		if !(aCodecs == nil || (f.NormACodec == "" && aCodecs.empty()) || aCodecs.member(f.NormACodec)) {
			continue
		}
		if !(vCodecs == nil || (f.NormVCodec == "" && vCodecs.empty()) || vCodecs.member(f.NormVCodec)) {
			continue
		}

		matched = append(matched, f)
	}

	sort.Sort(youtubedl.FormatByNormBR(matched))

	if len(matched) > 0 {
		return matched[0]
	}

	return nil
}

func findBestFormats(ydlFormats []*youtubedl.Format, format *Format) (aFormat *youtubedl.Format, vFormat *youtubedl.Format) {
	type neededFormat struct {
		aCodecs    *prioStringSet
		vCodecs    *prioStringSet
		aYDLFormat **youtubedl.Format
		vYDLFormat **youtubedl.Format
	}

	var neededFormats []*neededFormat

	// match exactly, both audio/video codecs found or not found
	neededFormats = append(neededFormats, &neededFormat{&format.ACodecs, &format.VCodecs, &aFormat, &vFormat})

	if !format.ACodecs.empty() {
		// matching audio codec with any video codec
		neededFormats = append(neededFormats, &neededFormat{&format.ACodecs, nil, &aFormat, nil})
		// match any audio codec and no video
		neededFormats = append(neededFormats, &neededFormat{nil, &prioStringSet{}, &aFormat, nil})
		// match any audio and video codec
		neededFormats = append(neededFormats, &neededFormat{nil, nil, &aFormat, nil})
	}
	if !format.VCodecs.empty() {
		// same logic as above
		neededFormats = append(neededFormats, &neededFormat{nil, &format.VCodecs, nil, &vFormat})
		neededFormats = append(neededFormats, &neededFormat{&prioStringSet{}, nil, nil, &vFormat})
		neededFormats = append(neededFormats, &neededFormat{nil, nil, nil, &vFormat})
	}

	// TODO: if only audio... choose stream with lowest video br?
	// TODO: rtmp only if mp4 etc?

	// prefer rtmp as it seems to send fragmented mp4 etc
	for _, proto := range []string{"rtmp", "*"} {
		for _, f := range neededFormats {
			m := findFormat(ydlFormats, proto, f.aCodecs, f.vCodecs)

			if m == nil {
				continue
			}

			if f.aYDLFormat != nil && *f.aYDLFormat == nil && m.NormACodec != "" {
				*f.aYDLFormat = m
			}
			if f.vYDLFormat != nil && *f.vYDLFormat == nil && m.NormVCodec != "" {
				*f.vYDLFormat = m
			}

			if (format.ACodecs.empty() || aFormat != nil) &&
				(format.VCodecs.empty() || vFormat != nil) {
				break
			}
		}
	}

	return aFormat, vFormat
}

func downloadAndProbeFormat(ydl *youtubedl.Info, filter string, debugLog *log.Logger) (r io.ReadCloser, pi *ffmpeg.ProbeInfo, err error) {
	var ydlStderr io.Writer
	if debugLog != nil {
		ydlStderr = &loggerWriter{Logger: debugLog, Prefix: "ydl 2> "}
	}
	r, err = ydl.Download(filter, ydlStderr)
	if err != nil {
		return nil, nil, err
	}

	rr := &reReadCloser{ReadCloser: r}

	var ffprobeStderr io.Writer
	if debugLog != nil {
		ffprobeStderr = &loggerWriter{Logger: debugLog, Prefix: fmt.Sprintf("ffprobe %s 2> ", filter)}
	}
	const maxProbeByteSize = 10 * 1024 * 1024
	pi, err = ffmpeg.FFprobe(io.LimitReader(rr, maxProbeByteSize), debugLog, ffprobeStderr)
	if err != nil {
		return nil, nil, err
	}
	// restart and replay buffer data used when probing
	rr.Restarted = true

	return rr, pi, nil
}

// YDLs youtubedl downloader with some extras
type YDLs struct {
	Formats Formats
}

// NewFromFile new YDLs using formats file
func NewFromFile(formatsPath string) (*YDLs, error) {
	formatsFile, err := os.Open(formatsPath)
	if err != nil {
		return nil, err
	}
	defer formatsFile.Close()
	formats, err := parseFormats(formatsFile)
	if err != nil {
		return nil, err
	}

	return &YDLs{Formats: formats}, nil
}

// Download downloads media from URL and makes sure output is in specified format
func (ydls *YDLs) Download(url string, formatName string, debugLog *log.Logger) (r io.ReadCloser, filename string, mimeType string, err error) {
	var closeOnDone []io.Closer
	closeOnDoneFn := func() {
		for _, c := range closeOnDone {
			c.Close()
		}
	}
	deferCloseFn := closeOnDoneFn
	defer func() {
		// will be nil if cmd starts and goroutine takes care of closing instead
		if deferCloseFn != nil {
			deferCloseFn()
		}
	}()

	log := log.New(ioutil.Discard, "", 0)
	if debugLog != nil {
		log = debugLog
	}

	log.Printf("URL: %s", url)

	var ydlStdout io.Writer
	if debugLog != nil {
		ydlStdout = &loggerWriter{Logger: debugLog, Prefix: "ydl 1> "}
	}
	ydl, err := youtubedl.NewFromURL(url, ydlStdout)
	if err != nil {
		log.Printf("Failed to download: %s", err)
		return nil, "", "", fmt.Errorf("failed to download")
	}

	log.Printf("Title: %s", ydl.Title)
	log.Printf("Available youtubedl formats:")
	for _, f := range ydl.Formats {
		log.Printf("  %s", f)
	}

	var outFormat *Format

	if formatName == "" {
		var probedInfo *ffmpeg.ProbeInfo
		r, probedInfo, err = downloadAndProbeFormat(ydl, "best[protocol=rtmp]/best", debugLog)
		if err != nil {
			return nil, "", "", err
		}

		log.Printf("Probed format %s", probedInfo)

		// see if we know about the probed format, otherwise fallback to "raw"
		outFormat = ydls.Formats.Find(probedInfo.FormatName(), probedInfo.ACodec(), probedInfo.VCodec())
	} else {
		outFormat = ydls.Formats.FindByName(formatName)
		if outFormat == nil {
			return nil, "", "", fmt.Errorf("could not find format")
		}

		aYDLFormat, vYDLFormat := findBestFormats(ydl.Formats, outFormat)

		log.Printf("Best youtubedl match for %s a=%s v=%s", formatName, aYDLFormat, vYDLFormat)

		var aProbedFormat *ffmpeg.ProbeInfo
		var aReader io.ReadCloser
		var aErr error
		var vProbedFormat *ffmpeg.ProbeInfo
		var vReader io.ReadCloser
		var vErr error

		// FIXME: messy

		if aYDLFormat != nil && vYDLFormat != nil {
			if aYDLFormat != vYDLFormat {
				var probeWG sync.WaitGroup
				probeWG.Add(2)
				go func() {
					defer probeWG.Done()
					aReader, aProbedFormat, aErr = downloadAndProbeFormat(ydl, aYDLFormat.FormatID, debugLog)
				}()
				go func() {
					defer probeWG.Done()
					vReader, vProbedFormat, vErr = downloadAndProbeFormat(ydl, vYDLFormat.FormatID, debugLog)
				}()
				probeWG.Wait()
				if aReader != nil {
					closeOnDone = append(closeOnDone, aReader)
				}
				if vReader != nil {
					closeOnDone = append(closeOnDone, vReader)
				}
			} else {
				aReader, aProbedFormat, aErr = downloadAndProbeFormat(ydl, aYDLFormat.FormatID, debugLog)
				vReader, vProbedFormat, vErr = aReader, aProbedFormat, aErr
				if aReader != nil {
					closeOnDone = append(closeOnDone, aReader)
				}
			}
		} else if aYDLFormat != nil && vYDLFormat == nil {
			aReader, aProbedFormat, aErr = downloadAndProbeFormat(ydl, aYDLFormat.FormatID, debugLog)
			if aReader != nil {
				closeOnDone = append(closeOnDone, aReader)
			}
		} else {
			aReader, aProbedFormat, aErr = downloadAndProbeFormat(ydl, "best", debugLog)
			vReader, vProbedFormat, vErr = aReader, aProbedFormat, aErr
			if aReader != nil {
				closeOnDone = append(closeOnDone, aReader)
			}
		}
		if err != nil {
			return nil, "", "", fmt.Errorf("failed to probe")
		}

		log.Printf("Stream mapping:")
		var maps []ffmpeg.Map
		if len(outFormat.ACodecs) > 0 && aProbedFormat != nil && aProbedFormat.ACodec() != "" {
			canCopy := outFormat.ACodecs.member(aProbedFormat.ACodec())
			ffmpegCodec := boolString(canCopy, "copy", outFormat.ACodecs.first())
			ydlACodec := "n/a"
			if aYDLFormat != nil {
				ydlACodec = aYDLFormat.NormACodec
			}

			maps = append(maps, ffmpeg.Map{
				Input:           aReader,
				Kind:            "audio",
				StreamSpecifier: "a:0",
				Codec:           ffmpegCodec,
				Flags:           outFormat.ACodecFlags,
			})
			log.Printf("  audio probed:%s ydl:%s -> %s", aProbedFormat.ACodec(), ydlACodec, ffmpegCodec)
		}
		if len(outFormat.VCodecs) > 0 && vProbedFormat != nil && vProbedFormat.VCodec() != "" {
			canCopy := outFormat.VCodecs.member(vProbedFormat.VCodec())
			ffmpegCodec := boolString(canCopy, "copy", outFormat.VCodecs.first())
			ydlVCodec := "n/a"
			if vYDLFormat != nil {
				ydlVCodec = vYDLFormat.NormVCodec
			}

			maps = append(maps, ffmpeg.Map{
				Input:           vReader,
				Kind:            "video",
				StreamSpecifier: "v:0",
				Codec:           ffmpegCodec,
				Flags:           outFormat.VCodecFlags,
			})
			log.Printf("  video probed:%s ydl:%s -> %s", vProbedFormat.VCodec(), ydlVCodec, ffmpegCodec)
		}

		var ffmpegStderr io.Writer
		if debugLog != nil {
			ffmpegStderr = &loggerWriter{Logger: debugLog, Prefix: "ffmpeg 2> "}
		}
		ffmpegR, ffmpegW := io.Pipe()
		closeOnDone = append(closeOnDone, ffmpegR)

		f := &ffmpeg.FFmpeg{
			Maps:     maps,
			Format:   ffmpeg.Format{Name: outFormat.Formats.first(), Flags: outFormat.FormatFlags},
			DebugLog: debugLog,
			Stdout:   ffmpegW,
			Stderr:   ffmpegStderr,
		}

		if err := f.Start(); err != nil {
			return nil, "", "", err
		}

		// probe read one byte to see if ffmpeg is happy
		b := make([]byte, 1)
		if _, err := io.ReadFull(ffmpegR, b); err != nil {
			if err := f.Wait(); err != nil {
				log.Printf("ffmpeg failed: %s", err)
				return nil, "", "", err
			}
			log.Printf("read failed: %s", err)
			return nil, "", "", err
		}

		// goroutine will take care of closing
		deferCloseFn = nil

		var w io.WriteCloser
		r, w = io.Pipe()
		closeOnDone = append(closeOnDone, w)

		go func() {
			defer closeOnDoneFn()

			if outFormat.Prepend == "id3v2" {
				writeID3v2FromYoutueDLInfo(w, ydl)
			}
			w.Write(b)
			io.Copy(w, ffmpegR)
			f.Wait()
		}()
	}

	if outFormat == nil {
		mimeType = "application/octet-stream"
		filename = ydl.Title + ".raw"
	} else {
		mimeType = outFormat.MIMEType
		filename = ydl.Title + "." + outFormat.Ext
	}

	return r, filename, mimeType, nil
}
<|endoftext|>


<commit_before>// This file was auto-generated using createmock. See the following page for
// more information:
//
//     https://github.com/jacobsa/oglemock
//

package mock_blob

import (
	blob "github.com/jacobsa/comeback/blob"
	fmt "fmt"
	oglemock "github.com/jacobsa/oglemock"
	runtime "runtime"
	unsafe "unsafe"
)

type MockStore interface {
	blob.Store
	oglemock.MockObject
}

type mockStore struct {
	controller	oglemock.Controller
	description	string
}

func NewMockStore(
	c oglemock.Controller,
	desc string) MockStore {
	return &mockStore{
		controller:	c,
		description:	desc,
	}
}

func (m *mockStore) Oglemock_Id() uintptr {
	return uintptr(unsafe.Pointer(m))
}

func (m *mockStore) Oglemock_Description() string {
	return m.description
}

func (m *mockStore) Load(p0 blob.Score) (o0 []uint8, o1 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"Load",
		file,
		line,
		[]interface{}{p0})

	if len(retVals) != 2 {
		panic(fmt.Sprintf("mockStore.Load: invalid return values: %v", retVals))
	}

	// o0 []uint8
	if retVals[0] != nil {
		o0 = retVals[0].([]uint8)
	}

	// o1 error
	if retVals[1] != nil {
		o1 = retVals[1].(error)
	}

	return
}

func (m *mockStore) Store(p0 []uint8) (o0 blob.Score, o1 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"Store",
		file,
		line,
		[]interface{}{p0})

	if len(retVals) != 2 {
		panic(fmt.Sprintf("mockStore.Store: invalid return values: %v", retVals))
	}

	// o0 blob.Score
	if retVals[0] != nil {
		o0 = retVals[0].(blob.Score)
	}

	// o1 error
	if retVals[1] != nil {
		o1 = retVals[1].(error)
	}

	return
}
<commit_msg>Fixed formatting.<commit_after>// This file was auto-generated using createmock. See the following page for
// more information:
//
//     https://github.com/jacobsa/oglemock
//

package mock_blob

import (
	fmt "fmt"
	blob "github.com/jacobsa/comeback/blob"
	oglemock "github.com/jacobsa/oglemock"
	runtime "runtime"
	unsafe "unsafe"
)

type MockStore interface {
	blob.Store
	oglemock.MockObject
}

type mockStore struct {
	controller  oglemock.Controller
	description string
}

func NewMockStore(
	c oglemock.Controller,
	desc string) MockStore {
	return &mockStore{
		controller:  c,
		description: desc,
	}
}

func (m *mockStore) Oglemock_Id() uintptr {
	return uintptr(unsafe.Pointer(m))
}

func (m *mockStore) Oglemock_Description() string {
	return m.description
}

func (m *mockStore) Load(p0 blob.Score) (o0 []uint8, o1 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"Load",
		file,
		line,
		[]interface{}{p0})

	if len(retVals) != 2 {
		panic(fmt.Sprintf("mockStore.Load: invalid return values: %v", retVals))
	}

	// o0 []uint8
	if retVals[0] != nil {
		o0 = retVals[0].([]uint8)
	}

	// o1 error
	if retVals[1] != nil {
		o1 = retVals[1].(error)
	}

	return
}

func (m *mockStore) Store(p0 []uint8) (o0 blob.Score, o1 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"Store",
		file,
		line,
		[]interface{}{p0})

	if len(retVals) != 2 {
		panic(fmt.Sprintf("mockStore.Store: invalid return values: %v", retVals))
	}

	// o0 blob.Score
	if retVals[0] != nil {
		o0 = retVals[0].(blob.Score)
	}

	// o1 error
	if retVals[1] != nil {
		o1 = retVals[1].(error)
	}

	return
}
<|endoftext|>


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

// This boilerplate code is based on proxiers in k8s.io/kubernetes/pkg/proxy to
// allow reuse of the rest of the proxy package without change

package proxy

import (
	"sync"
	"time"

	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/selection"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/record"
	k8sp "k8s.io/kubernetes/pkg/proxy"
	"k8s.io/kubernetes/pkg/proxy/apis"
	"k8s.io/kubernetes/pkg/proxy/config"
	"k8s.io/kubernetes/pkg/proxy/healthcheck"
	"k8s.io/kubernetes/pkg/util/async"
)

// Proxy watches for updates of Services and Endpoints, maintains their mapping
// and programs it into the dataplane
type Proxy interface {
	//Stop stops the proxy and waits for its exit
	Stop()
}

// DPSyncerState groups the information passed to the DPSyncer's Apply
type DPSyncerState struct {
	SvcMap       k8sp.ServiceMap
	EpsMap       k8sp.EndpointsMap
	StaleUDPEps  []k8sp.ServiceEndpoint
	StaleUDPSvcs sets.String
}

// DPSyncer is an interface representing the dataplane syncer that applies the
// observed changes to the dataplane
type DPSyncer interface {
	Apply(state DPSyncerState) error
	Stop()
}

type proxy struct {
	initState

	hostname string

	k8s kubernetes.Interface

	epsChanges *k8sp.EndpointChangeTracker
	svcChanges *k8sp.ServiceChangeTracker

	svcMap k8sp.ServiceMap
	epsMap k8sp.EndpointsMap

	dpSyncer DPSyncer
	// executes periodic the dataplane updates
	runner *async.BoundedFrequencyRunner
	// ensures that only one invocation runs at any time
	runnerLck sync.Mutex
	// sets the minimal distance between to sync to avoid overloading the
	// dataplane in case of frequent changes
	minDPSyncPeriod time.Duration

	// how often to fully sync with k8s - 0 is never
	syncPeriod time.Duration

	// event recorder to update node events
	recorder record.EventRecorder

	// LB health checker proxy
	healthChecker healthcheck.Server
	healthzServer healthcheck.HealthzUpdater

	stopCh   chan struct{}
	stopWg   sync.WaitGroup
	stopOnce sync.Once
}

// New returns a new Proxy for the given k8s interface
func New(k8s kubernetes.Interface, dp DPSyncer, hostname string, opts ...Option) (Proxy, error) {
	if k8s == nil {
		return nil, errors.Errorf("no k8s client")
	}

	if dp == nil {
		return nil, errors.Errorf("no dataplane syncer")
	}

	p := &proxy{
		k8s:      k8s,
		dpSyncer: dp,
		hostname: hostname,
		svcMap:   make(k8sp.ServiceMap),
		epsMap:   make(k8sp.EndpointsMap),

		recorder: new(loggerRecorder),

		minDPSyncPeriod: 30 * time.Second, // XXX revisit the default

		stopCh: make(chan struct{}),
	}

	for _, o := range opts {
		if err := o(p); err != nil {
			return nil, errors.WithMessage(err, "applying option")
		}
	}

	// We need to create the runner first as once we start getting updates, they
	// will kick it
	p.runner = async.NewBoundedFrequencyRunner("dp-sync-runner",
		p.invokeDPSyncer, p.minDPSyncPeriod, time.Hour /* XXX might be infinite? */, 1)
	p.startRoutine(func() { p.runner.Loop(p.stopCh) })

	p.healthChecker = healthcheck.NewServer(p.hostname, p.recorder, nil, nil)
	isIPv6 := false
	p.epsChanges = k8sp.NewEndpointChangeTracker(p.hostname,
		nil, // change if you want to provide more ctx
		&isIPv6,
		p.recorder,
		false, // endpointSlicesEnabled
	)
	p.svcChanges = k8sp.NewServiceChangeTracker(nil, &isIPv6, p.recorder)

	noProxyName, err := labels.NewRequirement(apis.LabelServiceProxyName, selection.DoesNotExist, nil)
	if err != nil {
		return nil, errors.Errorf("noProxyName selector: %s", err)
	}

	noHeadlessEndpoints, err := labels.NewRequirement(v1.IsHeadlessService, selection.DoesNotExist, nil)
	if err != nil {
		return nil, errors.Errorf("noHeadlessEndpoints selector: %s", err)
	}

	labelSelector := labels.NewSelector()
	labelSelector = labelSelector.Add(*noProxyName, *noHeadlessEndpoints)

	informerFactory := informers.NewSharedInformerFactoryWithOptions(k8s, p.syncPeriod,
		informers.WithTweakListOptions(func(options *metav1.ListOptions) {
			options.LabelSelector = labelSelector.String()
		}))

	svcConfig := config.NewServiceConfig(
		informerFactory.Core().V1().Services(),
		p.syncPeriod,
	)
	svcConfig.RegisterEventHandler(p)
	p.startRoutine(func() { svcConfig.Run(p.stopCh) })

	// TODO check if EndpointSlices are available and use them instead
	epsConfig := config.NewEndpointsConfig(informerFactory.Core().V1().Endpoints(), p.syncPeriod)
	epsConfig.RegisterEventHandler(p)
	p.startRoutine(func() { epsConfig.Run(p.stopCh) })

	p.startRoutine(func() { informerFactory.Start(p.stopCh) })

	return p, nil
}

func (p *proxy) Stop() {
	p.stopOnce.Do(func() {
		p.dpSyncer.Stop()
		close(p.stopCh)
		p.stopWg.Wait()
	})
}

func (p *proxy) startRoutine(f func()) {
	p.stopWg.Add(1)
	go func() {
		defer p.stopWg.Done()
		f()
	}()
}

func (p *proxy) syncDP() {
	p.runner.Run()
}

func (p *proxy) forceSyncDP() {
	p.invokeDPSyncer()
}

func (p *proxy) invokeDPSyncer() {
	if !p.isInitialized() {
		return
	}

	p.runnerLck.Lock()
	defer p.runnerLck.Unlock()

	svcUpdateResult := k8sp.UpdateServiceMap(p.svcMap, p.svcChanges)
	epsUpdateResult := p.epsMap.Update(p.epsChanges)

	staleUDPSvcs := svcUpdateResult.UDPStaleClusterIP

	/* XXX no tincluded in 1.15 yet
	// merge stale UDP services
	for _, svcPortName := range endpointUpdateResult.StaleServiceNames {
		if svcInfo, ok := p.svcMap[svcPortName]; ok && svcInfo != nil && svcInfo.Protocol() == v1.ProtocolUDP {
			staleUDPSvcs.Insert(svcInfo.ClusterIP().String())
			for _, extIP := range svcInfo.ExternalIPStrings() {
				staleUDPSvcs.Insert(extIP)
			}
		}
	}
	*/

	// XXX perhaps in a different thread that runs regularly
	if p.healthzServer != nil {
		p.healthzServer.UpdateTimestamp()
	}

	if err := p.healthChecker.SyncServices(svcUpdateResult.HCServiceNodePorts); err != nil {
		log.WithError(err).Error("Error syncing healthcheck services")
	}
	if err := p.healthChecker.SyncEndpoints(epsUpdateResult.HCEndpointsLocalIPSize); err != nil {
		log.WithError(err).Error("Error syncing healthcheck endpoints")
	}
	err := p.dpSyncer.Apply(DPSyncerState{
		SvcMap:       p.svcMap,
		EpsMap:       p.epsMap,
		StaleUDPSvcs: staleUDPSvcs,
	})

	if err != nil {
		log.WithError(err).Errorf("applying changes failed")
		// TODO log the error or panic as the best might be to restart
		// completely to wipe out the loaded bpf maps
	}
}

func (p *proxy) OnServiceAdd(svc *v1.Service) {
	p.OnServiceUpdate(nil, svc)
}

func (p *proxy) OnServiceUpdate(old, curr *v1.Service) {
	if p.svcChanges.Update(old, curr) && p.isInitialized() {
		p.syncDP()
	}
}

func (p *proxy) OnServiceDelete(svc *v1.Service) {
	p.OnServiceUpdate(svc, nil)
}

func (p *proxy) OnServiceSynced() {
	p.setSvcsSynced()
	p.forceSyncDP()
}

func (p *proxy) OnEndpointsAdd(eps *v1.Endpoints) {
	p.OnEndpointsUpdate(nil, eps)
}

func (p *proxy) OnEndpointsUpdate(old, curr *v1.Endpoints) {
	if p.epsChanges.Update(old, curr) && p.isInitialized() {
		p.syncDP()
	}
}

func (p *proxy) OnEndpointsDelete(eps *v1.Endpoints) {
	p.OnEndpointsUpdate(eps, nil)
}

func (p *proxy) OnEndpointsSynced() {
	p.setEpsSynced()
	p.forceSyncDP()
}

type initState struct {
	lck        sync.RWMutex
	svcsSynced bool
	epsSynced  bool
}

func (is *initState) isInitialized() bool {
	is.lck.RLock()
	defer is.lck.RUnlock()
	return is.svcsSynced && is.epsSynced
}

func (is *initState) setSvcsSynced() {
	is.lck.Lock()
	defer is.lck.Unlock()
	is.svcsSynced = true
}

func (is *initState) setEpsSynced() {
	is.lck.Lock()
	defer is.lck.Unlock()
	is.epsSynced = true
}

type loggerRecorder struct{}

func (r *loggerRecorder) Event(object runtime.Object, eventtype, reason, message string) {
}

func (r *loggerRecorder) Eventf(object runtime.Object, eventtype, reason,
	messageFmt string, args ...interface{}) {
}

func (r *loggerRecorder) PastEventf(object runtime.Object, timestamp metav1.Time, eventtype,
	reason, messageFmt string, args ...interface{}) {
}

func (r *loggerRecorder) AnnotatedEventf(object runtime.Object, annotations map[string]string,
	eventtype, reason, messageFmt string, args ...interface{}) {
}
<commit_msg>bpf: start kube-proxy threads after creating objects<commit_after>// Copyright (c) 2017-2020 Tigera, Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// This boilerplate code is based on proxiers in k8s.io/kubernetes/pkg/proxy to
// allow reuse of the rest of the proxy package without change

package proxy

import (
	"sync"
	"time"

	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/selection"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/record"
	k8sp "k8s.io/kubernetes/pkg/proxy"
	"k8s.io/kubernetes/pkg/proxy/apis"
	"k8s.io/kubernetes/pkg/proxy/config"
	"k8s.io/kubernetes/pkg/proxy/healthcheck"
	"k8s.io/kubernetes/pkg/util/async"
)

// Proxy watches for updates of Services and Endpoints, maintains their mapping
// and programs it into the dataplane
type Proxy interface {
	//Stop stops the proxy and waits for its exit
	Stop()
}

// DPSyncerState groups the information passed to the DPSyncer's Apply
type DPSyncerState struct {
	SvcMap       k8sp.ServiceMap
	EpsMap       k8sp.EndpointsMap
	StaleUDPEps  []k8sp.ServiceEndpoint
	StaleUDPSvcs sets.String
}

// DPSyncer is an interface representing the dataplane syncer that applies the
// observed changes to the dataplane
type DPSyncer interface {
	Apply(state DPSyncerState) error
	Stop()
}

type proxy struct {
	initState

	hostname string

	k8s kubernetes.Interface

	epsChanges *k8sp.EndpointChangeTracker
	svcChanges *k8sp.ServiceChangeTracker

	svcMap k8sp.ServiceMap
	epsMap k8sp.EndpointsMap

	dpSyncer DPSyncer
	// executes periodic the dataplane updates
	runner *async.BoundedFrequencyRunner
	// ensures that only one invocation runs at any time
	runnerLck sync.Mutex
	// sets the minimal distance between to sync to avoid overloading the
	// dataplane in case of frequent changes
	minDPSyncPeriod time.Duration

	// how often to fully sync with k8s - 0 is never
	syncPeriod time.Duration

	// event recorder to update node events
	recorder record.EventRecorder

	// LB health checker proxy
	healthChecker healthcheck.Server
	healthzServer healthcheck.HealthzUpdater

	stopCh   chan struct{}
	stopWg   sync.WaitGroup
	stopOnce sync.Once
}

// New returns a new Proxy for the given k8s interface
func New(k8s kubernetes.Interface, dp DPSyncer, hostname string, opts ...Option) (Proxy, error) {
	if k8s == nil {
		return nil, errors.Errorf("no k8s client")
	}

	if dp == nil {
		return nil, errors.Errorf("no dataplane syncer")
	}

	p := &proxy{
		k8s:      k8s,
		dpSyncer: dp,
		hostname: hostname,
		svcMap:   make(k8sp.ServiceMap),
		epsMap:   make(k8sp.EndpointsMap),

		recorder: new(loggerRecorder),

		minDPSyncPeriod: 30 * time.Second, // XXX revisit the default

		stopCh: make(chan struct{}),
	}

	for _, o := range opts {
		if err := o(p); err != nil {
			return nil, errors.WithMessage(err, "applying option")
		}
	}

	// We need to create the runner first as once we start getting updates, they
	// will kick it
	p.runner = async.NewBoundedFrequencyRunner("dp-sync-runner",
		p.invokeDPSyncer, p.minDPSyncPeriod, time.Hour /* XXX might be infinite? */, 1)

	p.healthChecker = healthcheck.NewServer(p.hostname, p.recorder, nil, nil)
	isIPv6 := false
	p.epsChanges = k8sp.NewEndpointChangeTracker(p.hostname,
		nil, // change if you want to provide more ctx
		&isIPv6,
		p.recorder,
		false, // endpointSlicesEnabled
	)
	p.svcChanges = k8sp.NewServiceChangeTracker(nil, &isIPv6, p.recorder)

	noProxyName, err := labels.NewRequirement(apis.LabelServiceProxyName, selection.DoesNotExist, nil)
	if err != nil {
		return nil, errors.Errorf("noProxyName selector: %s", err)
	}

	noHeadlessEndpoints, err := labels.NewRequirement(v1.IsHeadlessService, selection.DoesNotExist, nil)
	if err != nil {
		return nil, errors.Errorf("noHeadlessEndpoints selector: %s", err)
	}

	labelSelector := labels.NewSelector()
	labelSelector = labelSelector.Add(*noProxyName, *noHeadlessEndpoints)

	informerFactory := informers.NewSharedInformerFactoryWithOptions(k8s, p.syncPeriod,
		informers.WithTweakListOptions(func(options *metav1.ListOptions) {
			options.LabelSelector = labelSelector.String()
		}))

	svcConfig := config.NewServiceConfig(
		informerFactory.Core().V1().Services(),
		p.syncPeriod,
	)
	svcConfig.RegisterEventHandler(p)

	// TODO check if EndpointSlices are available and use them instead
	epsConfig := config.NewEndpointsConfig(informerFactory.Core().V1().Endpoints(), p.syncPeriod)
	epsConfig.RegisterEventHandler(p)

	p.startRoutine(func() { p.runner.Loop(p.stopCh) })
	p.startRoutine(func() { epsConfig.Run(p.stopCh) })
	p.startRoutine(func() { informerFactory.Start(p.stopCh) })
	p.startRoutine(func() { svcConfig.Run(p.stopCh) })

	return p, nil
}

func (p *proxy) Stop() {
	p.stopOnce.Do(func() {
		p.dpSyncer.Stop()
		close(p.stopCh)
		p.stopWg.Wait()
	})
}

func (p *proxy) startRoutine(f func()) {
	p.stopWg.Add(1)
	go func() {
		defer p.stopWg.Done()
		f()
	}()
}

func (p *proxy) syncDP() {
	p.runner.Run()
}

func (p *proxy) forceSyncDP() {
	p.invokeDPSyncer()
}

func (p *proxy) invokeDPSyncer() {
	if !p.isInitialized() {
		return
	}

	p.runnerLck.Lock()
	defer p.runnerLck.Unlock()

	svcUpdateResult := k8sp.UpdateServiceMap(p.svcMap, p.svcChanges)
	epsUpdateResult := p.epsMap.Update(p.epsChanges)

	staleUDPSvcs := svcUpdateResult.UDPStaleClusterIP

	/* XXX no tincluded in 1.15 yet
	// merge stale UDP services
	for _, svcPortName := range endpointUpdateResult.StaleServiceNames {
		if svcInfo, ok := p.svcMap[svcPortName]; ok && svcInfo != nil && svcInfo.Protocol() == v1.ProtocolUDP {
			staleUDPSvcs.Insert(svcInfo.ClusterIP().String())
			for _, extIP := range svcInfo.ExternalIPStrings() {
				staleUDPSvcs.Insert(extIP)
			}
		}
	}
	*/

	// XXX perhaps in a different thread that runs regularly
	if p.healthzServer != nil {
		p.healthzServer.UpdateTimestamp()
	}

	if err := p.healthChecker.SyncServices(svcUpdateResult.HCServiceNodePorts); err != nil {
		log.WithError(err).Error("Error syncing healthcheck services")
	}
	if err := p.healthChecker.SyncEndpoints(epsUpdateResult.HCEndpointsLocalIPSize); err != nil {
		log.WithError(err).Error("Error syncing healthcheck endpoints")
	}
	err := p.dpSyncer.Apply(DPSyncerState{
		SvcMap:       p.svcMap,
		EpsMap:       p.epsMap,
		StaleUDPSvcs: staleUDPSvcs,
	})

	if err != nil {
		log.WithError(err).Errorf("applying changes failed")
		// TODO log the error or panic as the best might be to restart
		// completely to wipe out the loaded bpf maps
	}
}

func (p *proxy) OnServiceAdd(svc *v1.Service) {
	p.OnServiceUpdate(nil, svc)
}

func (p *proxy) OnServiceUpdate(old, curr *v1.Service) {
	if p.svcChanges.Update(old, curr) && p.isInitialized() {
		p.syncDP()
	}
}

func (p *proxy) OnServiceDelete(svc *v1.Service) {
	p.OnServiceUpdate(svc, nil)
}

func (p *proxy) OnServiceSynced() {
	p.setSvcsSynced()
	p.forceSyncDP()
}

func (p *proxy) OnEndpointsAdd(eps *v1.Endpoints) {
	p.OnEndpointsUpdate(nil, eps)
}

func (p *proxy) OnEndpointsUpdate(old, curr *v1.Endpoints) {
	if p.epsChanges.Update(old, curr) && p.isInitialized() {
		p.syncDP()
	}
}

func (p *proxy) OnEndpointsDelete(eps *v1.Endpoints) {
	p.OnEndpointsUpdate(eps, nil)
}

func (p *proxy) OnEndpointsSynced() {
	p.setEpsSynced()
	p.forceSyncDP()
}

type initState struct {
	lck        sync.RWMutex
	svcsSynced bool
	epsSynced  bool
}

func (is *initState) isInitialized() bool {
	is.lck.RLock()
	defer is.lck.RUnlock()
	return is.svcsSynced && is.epsSynced
}

func (is *initState) setSvcsSynced() {
	is.lck.Lock()
	defer is.lck.Unlock()
	is.svcsSynced = true
}

func (is *initState) setEpsSynced() {
	is.lck.Lock()
	defer is.lck.Unlock()
	is.epsSynced = true
}

type loggerRecorder struct{}

func (r *loggerRecorder) Event(object runtime.Object, eventtype, reason, message string) {
}

func (r *loggerRecorder) Eventf(object runtime.Object, eventtype, reason,
	messageFmt string, args ...interface{}) {
}

func (r *loggerRecorder) PastEventf(object runtime.Object, timestamp metav1.Time, eventtype,
	reason, messageFmt string, args ...interface{}) {
}

func (r *loggerRecorder) AnnotatedEventf(object runtime.Object, annotations map[string]string,
	eventtype, reason, messageFmt string, args ...interface{}) {
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Clinton Freeman 2014
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
 * associated documentation files (the "Software"), to deal in the Software without restriction,
 * including without limitation the rights to use, copy, modify, merge, publish, distribute,
 * sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all copies or
 * substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
 * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

package main

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	// "github.com/gorilla/websocket"
	"reflect"
	"runtime"
	"testing"
	// "time"
)

func Test(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Message Suite")
}

var _ = Describe("Message", func() {
	Context("Utf8 encoding", func() {
		It("should return an error for non-utf8 encoded messages", func() {
			_, _, err := ParseMessage(string([]byte{0xff, 0xfe, 0xfd}))
			Ω(err).ShouldNot(BeNil())
		})

		It("should should not return an error for utf8 encoded messages", func() {
			_, _, err := ParseMessage("/announce|{\"id\":\"dc6ac0ae-6e15-409b-b211-228a8f4a43b9\"}|{\"browser\":\"node\",\"browserVersion\":\"?\",\"id\":\"dc6ac0ae-6e15-409b-b211-228a8f4a43b9\",\"agent\":\"signaller@0.18.3\",\"room\":\"test-room\"}")
			Ω(err).Should(BeNil())
		})
	})

	Context("Action parsing", func() {
		It("should be able to parse an announce message", func() {
			action, message, err := ParseMessage("/announce")
			Ω(err).Should(BeNil())
			Ω(runtime.FuncForPC(reflect.ValueOf(action).Pointer()).Name()).Should(Equal("github.com/cfreeman/signalbox.announce"))
			Ω(len(message)).Should(Equal(1))
		})

		It("should be able to parse a leave message", func() {
			action, message, err := ParseMessage("/leave")
			Ω(err).Should(BeNil())
			Ω(runtime.FuncForPC(reflect.ValueOf(action).Pointer()).Name()).Should(Equal("github.com/cfreeman/signalbox.leave"))
			Ω(len(message)).Should(Equal(1))
		})

		It("should be able to parse a to message", func() {
			action, message, err := ParseMessage("/to")
			Ω(err).Should(BeNil())
			Ω(runtime.FuncForPC(reflect.ValueOf(action).Pointer()).Name()).Should(Equal("github.com/cfreeman/signalbox.to"))
			Ω(len(message)).Should(Equal(1))
		})

		It("should be able to parse a custom message", func() {
			action, message, err := ParseMessage("/custom|part1|part2")
			Ω(err).Should(BeNil())
			Ω(runtime.FuncForPC(reflect.ValueOf(action).Pointer()).Name()).Should(Equal("github.com/cfreeman/signalbox.custom"))
			Ω(len(message)).Should(Equal(3))
			Ω(message[0]).Should(Equal("/custom"))
			Ω(message[1]).Should(Equal("part1"))
			Ω(message[2]).Should(Equal("part2"))
		})

		// TODO: Test Malformed messages.
	})

	Context("ParsePeerAndRoom", func() {
		It("should return an error when their is not enough parts to a message", func() {
			_, message, _ := ParseMessage("/foo")
			_, _, err := ParsePeerAndRoom(message)
			Ω(err).ShouldNot(BeNil())
		})

		It("should parse source id and room", func() {
			_, message, _ := ParseMessage("/announce|{\"id\":\"abc\"}|{\"room\":\"test\"}")
			source, destination, err := ParsePeerAndRoom(message)
			Ω(err).Should(BeNil())
			Ω(source.Id).Should(Equal("abc"))
			Ω(destination.Room).Should(Equal("test"))
		})
	})

	Context("Test SignalBox State", func() {
		var state SignalBox
		var announceAAct messageFn
		var announceAMsg []string

		var announceA2Act messageFn
		var announceA2Msg []string

		var leaveAAct messageFn
		var leaveAMsg []string

		var leaveA2Act messageFn
		var leaveA2Msg []string

		var announceBAct messageFn
		var announceBMsg []string

		var leaveBAct messageFn
		var leaveBMsg []string

		BeforeEach(func() {
			var err error
			state = SignalBox{make(map[string]*Peer),
				make(map[string]*Room),
				make(map[string]map[string]*Peer),
				make(map[string]map[string]*Room)}

			announceAAct, announceAMsg, err = ParseMessage("/announce|{\"id\":\"a\"}|{\"room\":\"test\"}")
			Ω(err).Should(BeNil())

			announceA2Act, announceA2Msg, err = ParseMessage("/announce|{\"id\":\"a\"}|{\"room\":\"test2\"}")
			Ω(err).Should(BeNil())

			leaveAAct, leaveAMsg, err = ParseMessage("/leave|{\"id\":\"a\"}|{\"room\":\"test\"}")
			Ω(err).Should(BeNil())

			leaveA2Act, leaveA2Msg, err = ParseMessage("/leave|{\"id\":\"a\"}|{\"room\":\"test2\"}")
			Ω(err).Should(BeNil())

			announceBAct, announceBMsg, err = ParseMessage("/announce|{\"id\":\"b\"}|{\"room\":\"test\"}")
			Ω(err).Should(BeNil())

			leaveBAct, leaveBMsg, err = ParseMessage("/leave|{\"id\":\"b\"}|{\"room\":\"test\"}")
			Ω(err).Should(BeNil())
		})

		It("only add someone to the roome once, even if they announce more than once", func() {
			state, err := announceAAct(announceAMsg, nil, state)
			Ω(err).Should(BeNil())
			state, err = announceAAct(announceAMsg, nil, state)
			Ω(err).Should(BeNil())

			Ω(len(state.Peers)).Should(Equal(1))
			Ω(len(state.Rooms)).Should(Equal(1))
			Ω(len(state.RoomContains)).Should(Equal(1))
			Ω(len(state.PeerIsIn)).Should(Equal(1))
			Ω(len(state.RoomContains["test"])).Should(Equal(1))
			Ω(len(state.PeerIsIn["a"])).Should(Equal(1))
			Ω(state.RoomContains["test"]["a"].Id).Should(Equal("a"))
			Ω(state.PeerIsIn["a"]["test"].Room).Should(Equal("test"))
		})

		It("should be able to add multiple people to a signalbox room", func() {
			state, err := announceAAct(announceAMsg, nil, state)
			Ω(err).Should(BeNil())
			state, err = announceBAct(announceBMsg, nil, state)
			Ω(err).Should(BeNil())

			Ω(len(state.Peers)).Should(Equal(2))
			Ω(len(state.Rooms)).Should(Equal(1))
			Ω(len(state.PeerIsIn)).Should(Equal(2))
			Ω(len(state.RoomContains)).Should(Equal(1))
			Ω(len(state.PeerIsIn["a"])).Should(Equal(1))
			Ω(len(state.PeerIsIn["b"])).Should(Equal(1))
			Ω(len(state.RoomContains["test"])).Should(Equal(2))
		})

		It("Should be able to have a person leave a signalbox room", func() {
			state, err := announceAAct(announceAMsg, nil, state)
			Ω(err).Should(BeNil())
			state, err = leaveAAct(leaveAMsg, nil, state)
			Ω(err).Should(BeNil())

			Ω(len(state.Peers)).Should(Equal(0))
			Ω(len(state.Rooms)).Should(Equal(0))
			Ω(len(state.PeerIsIn)).Should(Equal(0))
			Ω(len(state.RoomContains)).Should(Equal(0))
		})

		It("Should keep a room, if a person leaves but it still contains peers", func() {
			state, err := announceBAct(announceBMsg, nil, state)
			Ω(err).Should(BeNil())
			state, err = announceAAct(announceAMsg, nil, state)
			Ω(err).Should(BeNil())
			state, err = announceA2Act(announceA2Msg, nil, state)
			Ω(err).Should(BeNil())
			state, err = leaveBAct(leaveBMsg, nil, state)
			Ω(err).Should(BeNil())

			Ω(len(state.Peers)).Should(Equal(1))
			Ω(len(state.Rooms)).Should(Equal(2))
			Ω(len(state.PeerIsIn)).Should(Equal(1))
			Ω(len(state.RoomContains)).Should(Equal(2))
			Ω(len(state.PeerIsIn["a"])).Should(Equal(2))
			Ω(len(state.RoomContains["test"])).Should(Equal(1))
			Ω(len(state.RoomContains["test2"])).Should(Equal(1))
			Ω(state.RoomContains["test"]["a"].Id).Should(Equal("a"))
			Ω(state.RoomContains["test2"]["a"].Id).Should(Equal("a"))
		})
	})
})

// TODO: Port the rest of the tests over to Ginkgo.

// func connectPeer(id string, room string) (*websocket.Conn, error) {
// 	url := "ws://localhost:3000"
// 	res, _, err := websocket.DefaultDialer.Dial(url, nil)
// 	if err != nil || res == nil {
// 		return nil, err
// 	}

// 	connect := fmt.Sprintf("/announce|{\"id\":\"%s\"}|{\"room\":\"%s\"}", id, room)
// 	err = res.WriteMessage(websocket.TextMessage, []byte(connect))
// 	if err != nil {
// 		return nil, err
// 	}

// 	return res, nil
// }

// func TestAnnounceBroadcast(t *testing.T) {
// 	go main()

// 	a, err := connectPeer("a", "test-room")
// 	if err != nil {
// 		t.Error(err)
// 		return
// 	}

// 	b, err := connectPeer("b", "test-room")
// 	if err != nil {
// 		t.Error(err)
// 		return
// 	}

// 	_, message, err := a.ReadMessage()
// 	if err != nil || string(message) != "/announce|{\"id\":\"b\"}|{\"room\":\"test-room\"}" {
// 		t.Errorf("Peer A did not recieve the announce message for b.")
// 	}

// 	err = a.WriteMessage(websocket.TextMessage, []byte("/leave|{\"id\":\"a\"}|{\"room\":\"test-room\"}"))
// 	if err != nil {
// 		t.Error(err)
// 	}

// 	_, message, err = b.ReadMessage()
// 	if err != nil || string(message) != "/leave|{\"id\":\"a\"}|{\"room\":\"test-room\"}" {
// 		t.Errorf("Peer B did not recieve the leave message for a.")
// 	}
// }

// func TestMessage(t *testing.T) {
// 	a, err := connectPeer("a1", "test-to-message")
// 	if err != nil {
// 		t.Error(err)
// 		return
// 	}
// 	time.Sleep(2 * time.Millisecond)

// 	b, err := connectPeer("b2", "test-to-message")
// 	if err != nil {
// 		t.Error(err)
// 		return
// 	}
// 	time.Sleep(2 * time.Millisecond)

// 	c, err := connectPeer("c2", "test-to-message")
// 	if err != nil {
// 		t.Error(err)
// 		return
// 	}
// 	time.Sleep(2 * time.Millisecond)

// 	// Test custom Message.
// 	err = a.WriteMessage(websocket.TextMessage, []byte("/hello|{\"id\":\"a1\"}"))
// 	if err != nil {
// 		t.Error(err)
// 	}
// 	time.Sleep(2 * time.Millisecond)

// 	b.ReadMessage() // discard the c2 announce message.
// 	_, b_message, err := b.ReadMessage()
// 	if err != nil || string(b_message) != "/hello|{\"id\":\"a1\"}" {
// 		t.Errorf(string(b_message))
// 		t.Errorf("Peer B did not recieve the message from A.")
// 	}

// 	_, c_message, err := c.ReadMessage()
// 	if err != nil || string(c_message) != "/hello|{\"id\":\"a1\"}" {
// 		t.Errorf(string(c_message))
// 		t.Errorf("Peer C did not recieve the message from A.")
// 	}

// 	// Test TO Message.
// 	err = a.WriteMessage(websocket.TextMessage, []byte("/to|b2|/hello|{\"id\":\"a1\"}"))
// 	if err != nil {
// 		t.Error(err)
// 	}
// 	time.Sleep(2 * time.Millisecond)

// 	_, b_message, err = b.ReadMessage()
// 	if err != nil || string(b_message) != "/to|b2|/hello|{\"id\":\"a1\"}" {
// 		t.Errorf(string(b_message))
// 		t.Errorf("Peer B did not recieve the personal message from A.")
// 	}

// 	c.SetReadDeadline(time.Now().Add(4 * time.Millisecond))
// 	_, c_message, err = c.ReadMessage()
// 	if string(c_message) != "" {
// 		t.Errorf("Peer C was not expecting any messages.")
// 	}
// }
<commit_msg>Tested the parsing of malformed messages.<commit_after>/*
 * Copyright (c) Clinton Freeman 2014
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
 * associated documentation files (the "Software"), to deal in the Software without restriction,
 * including without limitation the rights to use, copy, modify, merge, publish, distribute,
 * sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all copies or
 * substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
 * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

package main

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	// "github.com/gorilla/websocket"
	"reflect"
	"runtime"
	"testing"
	// "time"
)

func Test(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Message Suite")
}

var _ = Describe("Message", func() {
	Context("Utf8 encoding", func() {
		It("should return an error for non-utf8 encoded messages", func() {
			_, _, err := ParseMessage(string([]byte{0xff, 0xfe, 0xfd}))
			Ω(err).ShouldNot(BeNil())
		})

		It("should should not return an error for utf8 encoded messages", func() {
			_, _, err := ParseMessage("/announce|{\"id\":\"dc6ac0ae-6e15-409b-b211-228a8f4a43b9\"}|{\"browser\":\"node\",\"browserVersion\":\"?\",\"id\":\"dc6ac0ae-6e15-409b-b211-228a8f4a43b9\",\"agent\":\"signaller@0.18.3\",\"room\":\"test-room\"}")
			Ω(err).Should(BeNil())
		})
	})

	Context("Action parsing", func() {
		It("should be able to parse an announce message", func() {
			action, message, err := ParseMessage("/announce")
			Ω(err).Should(BeNil())
			Ω(runtime.FuncForPC(reflect.ValueOf(action).Pointer()).Name()).Should(Equal("github.com/cfreeman/signalbox.announce"))
			Ω(len(message)).Should(Equal(1))
		})

		It("should be able to parse a leave message", func() {
			action, message, err := ParseMessage("/leave")
			Ω(err).Should(BeNil())
			Ω(runtime.FuncForPC(reflect.ValueOf(action).Pointer()).Name()).Should(Equal("github.com/cfreeman/signalbox.leave"))
			Ω(len(message)).Should(Equal(1))
		})

		It("should be able to parse a to message", func() {
			action, message, err := ParseMessage("/to")
			Ω(err).Should(BeNil())
			Ω(runtime.FuncForPC(reflect.ValueOf(action).Pointer()).Name()).Should(Equal("github.com/cfreeman/signalbox.to"))
			Ω(len(message)).Should(Equal(1))
		})

		It("should be able to parse a custom message", func() {
			action, message, err := ParseMessage("/custom|part1|part2")
			Ω(err).Should(BeNil())
			Ω(runtime.FuncForPC(reflect.ValueOf(action).Pointer()).Name()).Should(Equal("github.com/cfreeman/signalbox.custom"))
			Ω(len(message)).Should(Equal(3))
			Ω(message[0]).Should(Equal("/custom"))
			Ω(message[1]).Should(Equal("part1"))
			Ω(message[2]).Should(Equal("part2"))
		})

		It("should ignore malformed messages", func() {
			action, message, err := ParseMessage(":lkajsd??asdj/foo")
			Ω(err).Should(BeNil())
			Ω(len(message)).Should(Equal(1))
			Ω(runtime.FuncForPC(reflect.ValueOf(action).Pointer()).Name()).Should(Equal("github.com/cfreeman/signalbox.ignore"))
		})
	})

	Context("ParsePeerAndRoom", func() {
		It("should return an error when their is not enough parts to a message", func() {
			_, message, _ := ParseMessage("/foo")
			_, _, err := ParsePeerAndRoom(message)
			Ω(err).ShouldNot(BeNil())
		})

		It("should parse source id and room", func() {
			_, message, _ := ParseMessage("/announce|{\"id\":\"abc\"}|{\"room\":\"test\"}")
			source, destination, err := ParsePeerAndRoom(message)
			Ω(err).Should(BeNil())
			Ω(source.Id).Should(Equal("abc"))
			Ω(destination.Room).Should(Equal("test"))
		})
	})

	Context("Test SignalBox State", func() {
		var state SignalBox
		var announceAAct messageFn
		var announceAMsg []string

		var announceA2Act messageFn
		var announceA2Msg []string

		var leaveAAct messageFn
		var leaveAMsg []string

		var leaveA2Act messageFn
		var leaveA2Msg []string

		var announceBAct messageFn
		var announceBMsg []string

		var leaveBAct messageFn
		var leaveBMsg []string

		BeforeEach(func() {
			var err error
			state = SignalBox{make(map[string]*Peer),
				make(map[string]*Room),
				make(map[string]map[string]*Peer),
				make(map[string]map[string]*Room)}

			announceAAct, announceAMsg, err = ParseMessage("/announce|{\"id\":\"a\"}|{\"room\":\"test\"}")
			Ω(err).Should(BeNil())

			announceA2Act, announceA2Msg, err = ParseMessage("/announce|{\"id\":\"a\"}|{\"room\":\"test2\"}")
			Ω(err).Should(BeNil())

			leaveAAct, leaveAMsg, err = ParseMessage("/leave|{\"id\":\"a\"}|{\"room\":\"test\"}")
			Ω(err).Should(BeNil())

			leaveA2Act, leaveA2Msg, err = ParseMessage("/leave|{\"id\":\"a\"}|{\"room\":\"test2\"}")
			Ω(err).Should(BeNil())

			announceBAct, announceBMsg, err = ParseMessage("/announce|{\"id\":\"b\"}|{\"room\":\"test\"}")
			Ω(err).Should(BeNil())

			leaveBAct, leaveBMsg, err = ParseMessage("/leave|{\"id\":\"b\"}|{\"room\":\"test\"}")
			Ω(err).Should(BeNil())
		})

		It("only add someone to the roome once, even if they announce more than once", func() {
			state, err := announceAAct(announceAMsg, nil, state)
			Ω(err).Should(BeNil())
			state, err = announceAAct(announceAMsg, nil, state)
			Ω(err).Should(BeNil())

			Ω(len(state.Peers)).Should(Equal(1))
			Ω(len(state.Rooms)).Should(Equal(1))
			Ω(len(state.RoomContains)).Should(Equal(1))
			Ω(len(state.PeerIsIn)).Should(Equal(1))
			Ω(len(state.RoomContains["test"])).Should(Equal(1))
			Ω(len(state.PeerIsIn["a"])).Should(Equal(1))
			Ω(state.RoomContains["test"]["a"].Id).Should(Equal("a"))
			Ω(state.PeerIsIn["a"]["test"].Room).Should(Equal("test"))
		})

		It("should be able to add multiple people to a signalbox room", func() {
			state, err := announceAAct(announceAMsg, nil, state)
			Ω(err).Should(BeNil())
			state, err = announceBAct(announceBMsg, nil, state)
			Ω(err).Should(BeNil())

			Ω(len(state.Peers)).Should(Equal(2))
			Ω(len(state.Rooms)).Should(Equal(1))
			Ω(len(state.PeerIsIn)).Should(Equal(2))
			Ω(len(state.RoomContains)).Should(Equal(1))
			Ω(len(state.PeerIsIn["a"])).Should(Equal(1))
			Ω(len(state.PeerIsIn["b"])).Should(Equal(1))
			Ω(len(state.RoomContains["test"])).Should(Equal(2))
		})

		It("Should be able to have a person leave a signalbox room", func() {
			state, err := announceAAct(announceAMsg, nil, state)
			Ω(err).Should(BeNil())
			state, err = leaveAAct(leaveAMsg, nil, state)
			Ω(err).Should(BeNil())

			Ω(len(state.Peers)).Should(Equal(0))
			Ω(len(state.Rooms)).Should(Equal(0))
			Ω(len(state.PeerIsIn)).Should(Equal(0))
			Ω(len(state.RoomContains)).Should(Equal(0))
		})

		It("Should keep a room, if a person leaves but it still contains peers", func() {
			state, err := announceBAct(announceBMsg, nil, state)
			Ω(err).Should(BeNil())
			state, err = announceAAct(announceAMsg, nil, state)
			Ω(err).Should(BeNil())
			state, err = announceA2Act(announceA2Msg, nil, state)
			Ω(err).Should(BeNil())
			state, err = leaveBAct(leaveBMsg, nil, state)
			Ω(err).Should(BeNil())

			Ω(len(state.Peers)).Should(Equal(1))
			Ω(len(state.Rooms)).Should(Equal(2))
			Ω(len(state.PeerIsIn)).Should(Equal(1))
			Ω(len(state.RoomContains)).Should(Equal(2))
			Ω(len(state.PeerIsIn["a"])).Should(Equal(2))
			Ω(len(state.RoomContains["test"])).Should(Equal(1))
			Ω(len(state.RoomContains["test2"])).Should(Equal(1))
			Ω(state.RoomContains["test"]["a"].Id).Should(Equal("a"))
			Ω(state.RoomContains["test2"]["a"].Id).Should(Equal("a"))
		})
	})

	Context("Broadcast messages", func() {
		// func connectPeer(id string, room string) (*websocket.Conn, error) {
		// 	url := "ws://localhost:3000"
		// 	res, _, err := websocket.DefaultDialer.Dial(url, nil)
		// 	if err != nil || res == nil {
		// 		return nil, err
		// 	}

		// 	connect := fmt.Sprintf("/announce|{\"id\":\"%s\"}|{\"room\":\"%s\"}", id, room)
		// 	err = res.WriteMessage(websocket.TextMessage, []byte(connect))
		// 	if err != nil {
		// 		return nil, err
		// 	}

		// 	return res, nil
		// }

		go main()

	})
})

// TODO: Port the rest of the tests over to Ginkgo.

// func connectPeer(id string, room string) (*websocket.Conn, error) {
// 	url := "ws://localhost:3000"
// 	res, _, err := websocket.DefaultDialer.Dial(url, nil)
// 	if err != nil || res == nil {
// 		return nil, err
// 	}

// 	connect := fmt.Sprintf("/announce|{\"id\":\"%s\"}|{\"room\":\"%s\"}", id, room)
// 	err = res.WriteMessage(websocket.TextMessage, []byte(connect))
// 	if err != nil {
// 		return nil, err
// 	}

// 	return res, nil
// }

// func TestAnnounceBroadcast(t *testing.T) {
// 	go main()

// 	a, err := connectPeer("a", "test-room")
// 	if err != nil {
// 		t.Error(err)
// 		return
// 	}

// 	b, err := connectPeer("b", "test-room")
// 	if err != nil {
// 		t.Error(err)
// 		return
// 	}

// 	_, message, err := a.ReadMessage()
// 	if err != nil || string(message) != "/announce|{\"id\":\"b\"}|{\"room\":\"test-room\"}" {
// 		t.Errorf("Peer A did not recieve the announce message for b.")
// 	}

// 	err = a.WriteMessage(websocket.TextMessage, []byte("/leave|{\"id\":\"a\"}|{\"room\":\"test-room\"}"))
// 	if err != nil {
// 		t.Error(err)
// 	}

// 	_, message, err = b.ReadMessage()
// 	if err != nil || string(message) != "/leave|{\"id\":\"a\"}|{\"room\":\"test-room\"}" {
// 		t.Errorf("Peer B did not recieve the leave message for a.")
// 	}
// }

// func TestMessage(t *testing.T) {
// 	a, err := connectPeer("a1", "test-to-message")
// 	if err != nil {
// 		t.Error(err)
// 		return
// 	}
// 	time.Sleep(2 * time.Millisecond)

// 	b, err := connectPeer("b2", "test-to-message")
// 	if err != nil {
// 		t.Error(err)
// 		return
// 	}
// 	time.Sleep(2 * time.Millisecond)

// 	c, err := connectPeer("c2", "test-to-message")
// 	if err != nil {
// 		t.Error(err)
// 		return
// 	}
// 	time.Sleep(2 * time.Millisecond)

// 	// Test custom Message.
// 	err = a.WriteMessage(websocket.TextMessage, []byte("/hello|{\"id\":\"a1\"}"))
// 	if err != nil {
// 		t.Error(err)
// 	}
// 	time.Sleep(2 * time.Millisecond)

// 	b.ReadMessage() // discard the c2 announce message.
// 	_, b_message, err := b.ReadMessage()
// 	if err != nil || string(b_message) != "/hello|{\"id\":\"a1\"}" {
// 		t.Errorf(string(b_message))
// 		t.Errorf("Peer B did not recieve the message from A.")
// 	}

// 	_, c_message, err := c.ReadMessage()
// 	if err != nil || string(c_message) != "/hello|{\"id\":\"a1\"}" {
// 		t.Errorf(string(c_message))
// 		t.Errorf("Peer C did not recieve the message from A.")
// 	}

// 	// Test TO Message.
// 	err = a.WriteMessage(websocket.TextMessage, []byte("/to|b2|/hello|{\"id\":\"a1\"}"))
// 	if err != nil {
// 		t.Error(err)
// 	}
// 	time.Sleep(2 * time.Millisecond)

// 	_, b_message, err = b.ReadMessage()
// 	if err != nil || string(b_message) != "/to|b2|/hello|{\"id\":\"a1\"}" {
// 		t.Errorf(string(b_message))
// 		t.Errorf("Peer B did not recieve the personal message from A.")
// 	}

// 	c.SetReadDeadline(time.Now().Add(4 * time.Millisecond))
// 	_, c_message, err = c.ReadMessage()
// 	if string(c_message) != "" {
// 		t.Errorf("Peer C was not expecting any messages.")
// 	}
// }
<|endoftext|>


<commit_before>package meta

import (
	"errors"
	"io"
)

// VerifyPadding verifies that the padding metadata block only contains 0 bits.
// The provided io.Reader should limit the amount of data that can be read to
// header.Length bytes.
func VerifyPadding(r io.Reader) (err error) {
	// Verify up to 4 kb of padding each iteration.
	buf := make([]byte, 4096)
	for {
		n, err := r.Read(buf)
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		if !isAllZero(buf[:n]) {
			return errors.New("meta.VerifyPadding: invalid padding; must contain only zeroes")
		}
	}
	return nil
}

/// ### [ note ] ###
///    - Might trigger unnecessary errors.
/// ### [/ note ] ###

// isAllZero returns true if the value of each byte in the provided slice is 0,
// and false otherwise.
func isAllZero(buf []byte) bool {
	for _, b := range buf {
		if b != 0 {
			return false
		}
	}
	return true
}
<commit_msg>meta: Use array instead of slice in VerifyPadding.<commit_after>package meta

import (
	"errors"
	"io"
)

// VerifyPadding verifies that the padding metadata block only contains 0 bits.
// The provided io.Reader should limit the amount of data that can be read to
// header.Length bytes.
func VerifyPadding(r io.Reader) (err error) {
	// Verify up to 4 kb of padding each iteration.
	var buf [4096]byte
	for {
		n, err := r.Read(buf[:])
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		if !isAllZero(buf[:n]) {
			return errors.New("meta.VerifyPadding: invalid padding; must contain only zeroes")
		}
	}
	return nil
}

// isAllZero returns true if the value of each byte in the provided slice is 0,
// and false otherwise.
func isAllZero(buf []byte) bool {
	for _, b := range buf {
		if b != 0 {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package scipipe

import (
	"fmt"
	"os/exec"
	re "regexp"
	str "strings"
	// "time"
	"errors"
	"sync"
)

type ShellTask struct {
	task
	_OutOnly     bool
	InPorts      map[string]chan *FileTarget
	InPaths      map[string]string
	OutPorts     map[string]chan *FileTarget
	OutPathFuncs map[string]func() string
	ParamPorts   map[string]chan string
	Params       map[string]string
	Command      string
	Spawn        bool
}

func NewShellTask(command string) *ShellTask {
	return &ShellTask{
		Command:      command,
		InPorts:      make(map[string]chan *FileTarget),
		InPaths:      make(map[string]string),
		OutPorts:     make(map[string]chan *FileTarget),
		OutPathFuncs: make(map[string]func() string),
		ParamPorts:   make(map[string]chan string),
		Params:       make(map[string]string),
		Spawn:        true,
	}
}

func Sh(cmd string) *ShellTask {
	t := NewShellTask(cmd)
	t.initPortsFromCmdPattern(cmd)
	return t
}

func ShNoSpawn(cmd string) *ShellTask {
	t := NewShellTask(cmd)
	t.initPortsFromCmdPattern(cmd)
	t.Spawn = false
	return t
}

func ShParams(cmd string, params map[string]string) *ShellTask {
	t := NewShellTask(cmd)
	t.initPortsFromCmdPattern(cmd)
	if params != nil {
		// Send eternal list of options
		go func() {
			for name, val := range params {
				t.ParamPorts[name] <- val
			}
		}()
	}
	return t
}

func (t *ShellTask) initPortsFromCmdPattern(cmd string) {
	// Find in/out port names and Params and set up in struct fields
	r := getPlaceHolderRegex()
	ms := r.FindAllStringSubmatch(cmd, -1)
	for _, m := range ms {
		if len(m) < 3 {
			Check(errors.New("Too few matches"))
		}
		typ := m[1]
		name := m[2]
		if typ == "o" {
			t.OutPorts[name] = make(chan *FileTarget, BUFSIZE)
		} else if typ == "p" {
			t.ParamPorts[name] = make(chan string, BUFSIZE)
		}

		// else if typ == "i" {
		// Set up a channel on the inports, even though this is
		// often replaced by another tasks output port channel.
		// It might be nice to have it init'ed with a channel
		// anyways, for use cases when we want to send FileTargets
		// on the inport manually.
		// t.InPorts[name] = make(chan *FileTarget, BUFSIZE)
		// }
	}
}

func (t *ShellTask) Run() {
	// fmt.Println("Entering task: ", t.Command)
	defer t.closeOutChans()

	wg := new(sync.WaitGroup)
	mx := new(sync.Mutex)
	// Main loop
	for {
		inPortsClosed := t.receiveInputs()
		paramPortsClosed := t.receiveParams()

		if len(t.InPorts) == 0 && paramPortsClosed {
			// fmt.Println("Closing loop: No inports, and param ports closed")
			break
		}
		if len(t.ParamPorts) == 0 && inPortsClosed {
			// fmt.Println("Closing loop: No inports, and in ports closed")
			break
		}
		if inPortsClosed && paramPortsClosed {
			// fmt.Println("Closing loop: Both inports and param ports closed")
			break
		}

		// Really needed?
		outPaths := copyMapStrStr(t.createOutPaths())

		// Format
		cmd := t.formatCommand(t.Command)

		if t.Spawn {
			wg.Add(1)
			go func() {
				t.executeCommand(cmd)
				t.sendOutputs(outPaths, mx)
				wg.Done()
			}()
		} else {
			t.executeCommand(cmd)
			t.sendOutputs(outPaths, mx)
		}

		// If there are no inports, we know we should exit the loop
		// directly after executing the command, and sending the outputs
		if len(t.InPorts) == 0 && len(t.ParamPorts) == 0 {
			fmt.Println("Closing after send: No inports or param ports")
			break
		}
	}
	fmt.Printf("Starting to wait (task '%s')\n", t.Command)
	wg.Wait()
	fmt.Printf("Finished waiting (task '%s')\n", t.Command)
	// fmt.Println("Exiting task:  ", t.Command)
}

func (t *ShellTask) closeOutChans() {
	// Close output channels
	for _, ochan := range t.OutPorts {
		close(ochan)
	}
}

func (t *ShellTask) receiveInputs() bool {
	inPortsClosed := false
	// Read input targets on in-ports and set up path mappings
	for iname, ichan := range t.InPorts {
		infile, open := <-ichan
		if !open {
			inPortsClosed = true
			continue
		}
		// fmt.Println("Receiving file:", infile.GetPath())
		t.InPaths[iname] = infile.GetPath()
	}
	return inPortsClosed
}

func (t *ShellTask) receiveParams() bool {
	paramPortsClosed := false
	// Read input targets on in-ports and set up path mappings
	for pname, pchan := range t.ParamPorts {
		pval, open := <-pchan
		if !open {
			paramPortsClosed = true
			continue
		}
		// fmt.Println("Receiving param:", pname, "with value", pval)
		t.Params[pname] = pval
	}
	return paramPortsClosed
}

func (t *ShellTask) sendOutputs(outPaths map[string]string, mx *sync.Mutex) {
	// Send output targets on out ports
	mx.Lock()
	for oname, ochan := range t.OutPorts {
		outName := outPaths[oname]
		ft := NewFileTarget(outName)
		// fmt.Println("Sending file:  ", ft.GetPath())
		ochan <- ft
	}
	mx.Unlock()
}

func (t *ShellTask) createOutPaths() (outPaths map[string]string) {
	outPaths = make(map[string]string)
	for oname, ofun := range t.OutPathFuncs {
		outPaths[oname] = ofun()
	}
	return outPaths
}

func (t *ShellTask) executeCommand(cmd string) {
	fmt.Println("Executing cmd: ", cmd)
	_, err := exec.Command("bash", "-c", cmd).Output()
	// fmt.Println("Command output: ", string(cmdOut))
	Check(err)
}

func (t *ShellTask) formatCommand(cmd string) string {
	r := getPlaceHolderRegex()
	ms := r.FindAllStringSubmatch(cmd, -1)
	for _, m := range ms {
		whole := m[0]
		typ := m[1]
		name := m[2]
		var newstr string
		if typ == "o" {
			if t.OutPathFuncs[name] == nil {
				msg := fmt.Sprint("Missing outpath function for outport '", name, "' of shell task '", t.Command, "'")
				Check(errors.New(msg))
			} else {
				newstr = t.OutPathFuncs[name]()
			}
		} else if typ == "i" {
			if t.InPaths[name] == "" {
				msg := fmt.Sprint("Missing inpath for inport '", name, "' of shell task '", t.Command, "'")
				Check(errors.New(msg))
			} else {
				newstr = t.InPaths[name]
			}
		} else if typ == "p" {
			if t.Params[name] == "" {
				msg := fmt.Sprint("Missing param value param '", name, "' of shell task '", t.Command, "'")
				Check(errors.New(msg))
			} else {
				newstr = t.Params[name]
			}
		}
		if newstr == "" {
			msg := fmt.Sprint("Replace failed for port ", name, " in task '", t.Command, "'")
			Check(errors.New(msg))
		}
		cmd = str.Replace(cmd, whole, newstr, -1)
	}
	return cmd
}

func (t *ShellTask) GetInPath(inPort string) string {
	var inPath string
	if t.InPaths[inPort] != "" {
		inPath = t.InPaths[inPort]
	} else {
		msg := fmt.Sprint("t.GetInPath(): Missing inpath for inport '", inPort, "' of shell task '", t.Command, "'")
		Check(errors.New(msg))
	}
	return inPath
}

func getPlaceHolderRegex() *re.Regexp {
	r, err := re.Compile("{(o|i|p):([^{}:]+)}")
	Check(err)
	return r
}
<commit_msg>Enabling to set static parameters<commit_after>package scipipe

import (
	"fmt"
	"os/exec"
	re "regexp"
	str "strings"
	// "time"
	"errors"
	"sync"
)

type ShellTask struct {
	task
	_OutOnly     bool
	InPorts      map[string]chan *FileTarget
	InPaths      map[string]string
	OutPorts     map[string]chan *FileTarget
	OutPathFuncs map[string]func() string
	ParamPorts   map[string]chan string
	Params       map[string]string
	Command      string
	Spawn        bool
}

func NewShellTask(command string) *ShellTask {
	return &ShellTask{
		Command:      command,
		InPorts:      make(map[string]chan *FileTarget),
		InPaths:      make(map[string]string),
		OutPorts:     make(map[string]chan *FileTarget),
		OutPathFuncs: make(map[string]func() string),
		ParamPorts:   make(map[string]chan string),
		Params:       make(map[string]string),
		Spawn:        true,
	}
}

func Sh(cmd string) *ShellTask {
	t := NewShellTask(cmd)
	t.initPortsFromCmdPattern(cmd, nil)
	return t
}

func ShParams(cmd string, params map[string]string) *ShellTask {
	t := NewShellTask(cmd)
	t.initPortsFromCmdPattern(cmd, params)
	return t
}

func (t *ShellTask) initPortsFromCmdPattern(cmd string, params map[string]string) {
	// Find in/out port names and Params and set up in struct fields
	r := getPlaceHolderRegex()
	ms := r.FindAllStringSubmatch(cmd, -1)
	for _, m := range ms {
		if len(m) < 3 {
			Check(errors.New("Too few matches"))
		}
		typ := m[1]
		name := m[2]
		if typ == "o" {
			t.OutPorts[name] = make(chan *FileTarget, BUFSIZE)
		} else if typ == "p" {
			if params == nil {
				t.ParamPorts[name] = make(chan string, BUFSIZE)
			} else {
				t.Params[name] = params[name]
			}
		}

		// else if typ == "i" {
		// Set up a channel on the inports, even though this is
		// often replaced by another tasks output port channel.
		// It might be nice to have it init'ed with a channel
		// anyways, for use cases when we want to send FileTargets
		// on the inport manually.
		// t.InPorts[name] = make(chan *FileTarget, BUFSIZE)
		// }
	}
}

func (t *ShellTask) Run() {
	// fmt.Println("Entering task: ", t.Command)
	defer t.closeOutChans()

	wg := new(sync.WaitGroup)
	mx := new(sync.Mutex)
	// Main loop
	for {
		inPortsClosed := t.receiveInputs()
		paramPortsClosed := t.receiveParams()

		if len(t.InPorts) == 0 && paramPortsClosed {
			// fmt.Println("Closing loop: No inports, and param ports closed")
			break
		}
		if len(t.ParamPorts) == 0 && inPortsClosed {
			// fmt.Println("Closing loop: No inports, and in ports closed")
			break
		}
		if inPortsClosed && paramPortsClosed {
			// fmt.Println("Closing loop: Both inports and param ports closed")
			break
		}

		// Really needed?
		outPaths := copyMapStrStr(t.createOutPaths())

		// Format
		cmd := t.formatCommand(t.Command)

		if t.Spawn {
			wg.Add(1)
			go func() {
				t.executeCommand(cmd)
				t.sendOutputs(outPaths, mx)
				wg.Done()
			}()
		} else {
			t.executeCommand(cmd)
			t.sendOutputs(outPaths, mx)
		}

		// If there are no inports, we know we should exit the loop
		// directly after executing the command, and sending the outputs
		if len(t.InPorts) == 0 && len(t.ParamPorts) == 0 {
			fmt.Println("Closing after send: No inports or param ports")
			break
		}
	}
	fmt.Printf("Starting to wait (task '%s')\n", t.Command)
	wg.Wait()
	fmt.Printf("Finished waiting (task '%s')\n", t.Command)
	// fmt.Println("Exiting task:  ", t.Command)
}

func (t *ShellTask) closeOutChans() {
	// Close output channels
	for _, ochan := range t.OutPorts {
		close(ochan)
	}
}

func (t *ShellTask) receiveInputs() bool {
	inPortsClosed := false
	// Read input targets on in-ports and set up path mappings
	for iname, ichan := range t.InPorts {
		infile, open := <-ichan
		if !open {
			inPortsClosed = true
			continue
		}
		// fmt.Println("Receiving file:", infile.GetPath())
		t.InPaths[iname] = infile.GetPath()
	}
	return inPortsClosed
}

func (t *ShellTask) receiveParams() bool {
	paramPortsClosed := false
	// Read input targets on in-ports and set up path mappings
	for pname, pchan := range t.ParamPorts {
		pval, open := <-pchan
		if !open {
			paramPortsClosed = true
			continue
		}
		// fmt.Println("Receiving param:", pname, "with value", pval)
		t.Params[pname] = pval
	}
	return paramPortsClosed
}

func (t *ShellTask) sendOutputs(outPaths map[string]string, mx *sync.Mutex) {
	// Send output targets on out ports
	mx.Lock()
	for oname, ochan := range t.OutPorts {
		outName := outPaths[oname]
		ft := NewFileTarget(outName)
		// fmt.Println("Sending file:  ", ft.GetPath())
		ochan <- ft
	}
	mx.Unlock()
}

func (t *ShellTask) createOutPaths() (outPaths map[string]string) {
	outPaths = make(map[string]string)
	for oname, ofun := range t.OutPathFuncs {
		outPaths[oname] = ofun()
	}
	return outPaths
}

func (t *ShellTask) executeCommand(cmd string) {
	fmt.Println("Executing cmd: ", cmd)
	_, err := exec.Command("bash", "-c", cmd).Output()
	// fmt.Println("Command output: ", string(cmdOut))
	Check(err)
}

func (t *ShellTask) formatCommand(cmd string) string {
	r := getPlaceHolderRegex()
	ms := r.FindAllStringSubmatch(cmd, -1)
	for _, m := range ms {
		whole := m[0]
		typ := m[1]
		name := m[2]
		var newstr string
		if typ == "o" {
			if t.OutPathFuncs[name] == nil {
				msg := fmt.Sprint("Missing outpath function for outport '", name, "' of shell task '", t.Command, "'")
				Check(errors.New(msg))
			} else {
				newstr = t.OutPathFuncs[name]()
			}
		} else if typ == "i" {
			if t.InPaths[name] == "" {
				msg := fmt.Sprint("Missing inpath for inport '", name, "' of shell task '", t.Command, "'")
				Check(errors.New(msg))
			} else {
				newstr = t.InPaths[name]
			}
		} else if typ == "p" {
			if t.Params[name] == "" {
				msg := fmt.Sprint("Missing param value param '", name, "' of shell task '", t.Command, "'")
				Check(errors.New(msg))
			} else {
				newstr = t.Params[name]
			}
		}
		if newstr == "" {
			msg := fmt.Sprint("Replace failed for port ", name, " in task '", t.Command, "'")
			Check(errors.New(msg))
		}
		cmd = str.Replace(cmd, whole, newstr, -1)
	}
	return cmd
}

func (t *ShellTask) GetInPath(inPort string) string {
	var inPath string
	if t.InPaths[inPort] != "" {
		inPath = t.InPaths[inPort]
	} else {
		msg := fmt.Sprint("t.GetInPath(): Missing inpath for inport '", inPort, "' of shell task '", t.Command, "'")
		Check(errors.New(msg))
	}
	return inPath
}

func getPlaceHolderRegex() *re.Regexp {
	r, err := re.Compile("{(o|i|p):([^{}:]+)}")
	Check(err)
	return r
}
<|endoftext|>


<commit_before>// SMTP server package.
package smtpd

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"net"
	"os"
	"regexp"
	"strings"
	"time"
)

var (
	rcptToRE    = regexp.MustCompile(`[Tt][Oo]:<(.+)>`)
	mailFromRE  = regexp.MustCompile(`[Ff][Rr][Oo][Mm]:<(.*)>`) // Delivery Status Notifications are sent with "MAIL FROM:<>"
)

// Definition of handler function.
type Handler func(remoteAddr net.Addr, from string, to []string, data []byte)

// ListenAndServe listens on the TCP network address addr
// and then calls Serve with handler to handle requests
// on incoming connections.
func ListenAndServe(addr string, handler Handler, appname string, hostname string) error {
	srv := &Server{Addr: addr, Handler: handler, Appname: appname, Hostname: hostname}
	return srv.ListenAndServe()
}

// Server is an SMTP server.
type Server struct {
	Addr    string // TCP address to listen on, defaults to ":25" (all addresses, port 25) if empty
	Handler Handler
	Appname string
	Hostname string
}

// ListenAndServe listens on the TCP network address srv.Addr and then
// calls Serve to handle requests on incoming connections.  If
// srv.Addr is blank, ":25" is used.
func (srv *Server) ListenAndServe() error {
	if srv.Addr == "" {
		srv.Addr = ":25"
	}
	if srv.Appname == "" {
		srv.Appname = "smtpd"
	}
	if srv.Hostname == "" {
		srv.Hostname, _ = os.Hostname()
	}
	ln, err := net.Listen("tcp", srv.Addr)
	if err != nil {
		return err
	}
	return srv.Serve(ln)
}

func (srv *Server) Serve(ln net.Listener) error {
	defer ln.Close()
	for {
		conn, err := ln.Accept()
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
				log.Printf("%s: Accept error: %v", srv.Appname, err)
				continue
			}
			return err
		}
		session, err := srv.newSession(conn)
		if err != nil {
			continue
		}
		go session.serve()
	}
}

type session struct {
	srv        *Server
	conn       net.Conn
	br         *bufio.Reader
	bw         *bufio.Writer
	remoteIP   string // Remote IP address
	remoteHost string // Remote hostname according to reverse DNS lookup
	remoteName string // Remote hostname as supplied with EHLO
}

// Create new session from connection.
func (srv *Server) newSession(conn net.Conn) (s *session, err error) {
	s = &session{
		srv:  srv,
		conn: conn,
		br:   bufio.NewReader(conn),
		bw:   bufio.NewWriter(conn),
	}
	return
}

// Function called to handle connection requests.
func (s *session) serve() {
	defer s.conn.Close()
	var from string
	var to []string
	var buffer bytes.Buffer

	// Get remote end info for the Received header.
	s.remoteIP, _, _ = net.SplitHostPort(s.conn.RemoteAddr().String())
	names, err := net.LookupAddr(s.remoteIP)
	if err == nil && len(names) > 0 {
		s.remoteHost = names[0]
	} else {
		s.remoteHost = "unknown"
	}

	// Send banner.
	s.writef("220 %s %s SMTP Service ready", s.srv.Hostname, s.srv.Appname)

loop:
	for {
		// Attempt to read a line from the socket.
		// On error, assume the client has gone away i.e. return from serve().
		line, err := s.readLine()
		if err != nil {
			break
		}
		verb, args := s.parseLine(line)

		switch verb {
		case "EHLO", "HELO":
			s.remoteName = args
			s.writef("250 %s greets %s", s.srv.Hostname, s.remoteName)

			// RFC 2821 section 4.1.4 specifies that EHLO has the same effect as RSET.
			from = ""
			to = nil
			buffer.Reset()
		case "MAIL":
			match := mailFromRE.FindStringSubmatch(args)
			if match == nil {
				s.writef("501 Syntax error in parameters or arguments (invalid FROM parameter)")
			} else {
				from = match[1]
				s.writef("250 Ok")
			}
			to = nil
			buffer.Reset()
		case "RCPT":
			if from == "" {
				s.writef("503 Bad sequence of commands (MAIL required before RCPT)")
				break
			}

			match := rcptToRE.FindStringSubmatch(args)
			if match == nil {
				s.writef("501 Syntax error in parameters or arguments (invalid TO parameter)")
			} else {
				// RFC 5321 specifies 100 minimum recipients
				if len(to) == 100 {
					s.writef("452 Too many recipients")
				} else {
					to = append(to, match[1])
					s.writef("250 Ok")
				}
			}
		case "DATA":
			if from == "" || to == nil {
				s.writef("503 Bad sequence of commands (MAIL & RCPT required before DATA)")
				break
			}

			s.writef("354 Start mail input; end with <CR><LF>.<CR><LF>")

			// Attempt to read message body from the socket.
			// On error, assume the client has gone away i.e. return from serve().
			data, err := s.readData()
			if err != nil {
				break loop
			}

			// Create Received header & write message body into buffer.
			buffer.Reset()
			buffer.Write(s.makeHeaders(to))
			buffer.Write(data)
			s.writef("250 Ok: queued")

			// Pass mail on to handler.
			if s.srv.Handler != nil {
				go s.srv.Handler(s.conn.RemoteAddr(), from, to, buffer.Bytes())
			}

			// Reset for next mail.
			from = ""
			to = nil
			buffer.Reset()
		case "QUIT":
			s.writef("221 %s %s SMTP Service closing transmission channel", s.srv.Hostname, s.srv.Appname)
			break loop
		case "RSET":
			s.writef("250 Ok")
			from = ""
			to = nil
			buffer.Reset()
		case "NOOP":
			s.writef("250 Ok")
		case "HELP", "VRFY", "EXPN":
			// See RFC 5321 section 4.2.4 for usage of 500 & 502 reply codes
			s.writef("502 Command not implemented")
		default:
			// See RFC 5321 section 4.2.4 for usage of 500 & 502 reply codes
			s.writef("500 Syntax error, command unrecognized")
		}
	}
}

// Wrapper function for writing a complete line to the socket.
func (s *session) writef(format string, args ...interface{}) {
	fmt.Fprintf(s.bw, format+"\r\n", args...)
	s.bw.Flush()
}

// Read a complete line from the socket.
func (s *session) readLine() (string, error) {
	line, err := s.br.ReadString('\n')
	if err != nil {
		return "", err
	}
	line = strings.TrimSpace(line) // Strip trailing \r\n
	return line, err
}

// Parse a line read from the socket.
func (s *session) parseLine(line string) (verb string, args string) {
	if idx := strings.Index(line, " "); idx != -1 {
		verb = strings.ToUpper(line[:idx])
		args = strings.TrimSpace(line[idx+1 : len(line)])
	} else {
		verb = strings.ToUpper(line)
		args = ""
	}
	return verb, args
}

// Read the message data following a DATA command.
func (s *session) readData() ([]byte, error) {
	var data []byte
	for {
		slice, err := s.br.ReadSlice('\n')
		if err != nil {
			return nil, err
		}
		// Handle end of data denoted by lone period (\r\n.\r\n)
		if bytes.Equal(slice, []byte(".\r\n")) {
			break
		}
		// Remove leading period (RFC 5321 section 4.5.2)
		if slice[0] == '.' {
			slice = slice[1:]
		}
		data = append(data, slice...)
	}
	return data, nil
}

// Create the Received header to comply with RFC 2821 section 3.8.2.
// TODO: Work out what to do with multiple to addresses.
func (s *session) makeHeaders(to []string) []byte {
	var buffer bytes.Buffer
	now := time.Now().Format("Mon, _2 Jan 2006 15:04:05 -0700 (MST)")
	buffer.WriteString(fmt.Sprintf("Received: from %s (%s [%s])\r\n", s.remoteName, s.remoteHost, s.remoteIP))
	buffer.WriteString(fmt.Sprintf("        by %s (%s) with SMTP\r\n", s.srv.Hostname, s.srv.Appname))
	buffer.WriteString(fmt.Sprintf("        for <%s>; %s\r\n", to[0], now))
	return buffer.Bytes()
}
<commit_msg>Formatting cleanup<commit_after>// SMTP server package.
package smtpd

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"net"
	"os"
	"regexp"
	"strings"
	"time"
)

var (
	rcptToRE   = regexp.MustCompile(`[Tt][Oo]:<(.+)>`)
	mailFromRE = regexp.MustCompile(`[Ff][Rr][Oo][Mm]:<(.*)>`) // Delivery Status Notifications are sent with "MAIL FROM:<>"
)

// Definition of handler function.
type Handler func(remoteAddr net.Addr, from string, to []string, data []byte)

// ListenAndServe listens on the TCP network address addr
// and then calls Serve with handler to handle requests
// on incoming connections.
func ListenAndServe(addr string, handler Handler, appname string, hostname string) error {
	srv := &Server{Addr: addr, Handler: handler, Appname: appname, Hostname: hostname}
	return srv.ListenAndServe()
}

// Server is an SMTP server.
type Server struct {
	Addr     string // TCP address to listen on, defaults to ":25" (all addresses, port 25) if empty
	Handler  Handler
	Appname  string
	Hostname string
}

// ListenAndServe listens on the TCP network address srv.Addr and then
// calls Serve to handle requests on incoming connections.  If
// srv.Addr is blank, ":25" is used.
func (srv *Server) ListenAndServe() error {
	if srv.Addr == "" {
		srv.Addr = ":25"
	}
	if srv.Appname == "" {
		srv.Appname = "smtpd"
	}
	if srv.Hostname == "" {
		srv.Hostname, _ = os.Hostname()
	}
	ln, err := net.Listen("tcp", srv.Addr)
	if err != nil {
		return err
	}
	return srv.Serve(ln)
}

func (srv *Server) Serve(ln net.Listener) error {
	defer ln.Close()
	for {
		conn, err := ln.Accept()
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
				log.Printf("%s: Accept error: %v", srv.Appname, err)
				continue
			}
			return err
		}
		session, err := srv.newSession(conn)
		if err != nil {
			continue
		}
		go session.serve()
	}
}

type session struct {
	srv        *Server
	conn       net.Conn
	br         *bufio.Reader
	bw         *bufio.Writer
	remoteIP   string // Remote IP address
	remoteHost string // Remote hostname according to reverse DNS lookup
	remoteName string // Remote hostname as supplied with EHLO
}

// Create new session from connection.
func (srv *Server) newSession(conn net.Conn) (s *session, err error) {
	s = &session{
		srv:  srv,
		conn: conn,
		br:   bufio.NewReader(conn),
		bw:   bufio.NewWriter(conn),
	}
	return
}

// Function called to handle connection requests.
func (s *session) serve() {
	defer s.conn.Close()
	var from string
	var to []string
	var buffer bytes.Buffer

	// Get remote end info for the Received header.
	s.remoteIP, _, _ = net.SplitHostPort(s.conn.RemoteAddr().String())
	names, err := net.LookupAddr(s.remoteIP)
	if err == nil && len(names) > 0 {
		s.remoteHost = names[0]
	} else {
		s.remoteHost = "unknown"
	}

	// Send banner.
	s.writef("220 %s %s SMTP Service ready", s.srv.Hostname, s.srv.Appname)

loop:
	for {
		// Attempt to read a line from the socket.
		// On error, assume the client has gone away i.e. return from serve().
		line, err := s.readLine()
		if err != nil {
			break
		}
		verb, args := s.parseLine(line)

		switch verb {
		case "EHLO", "HELO":
			s.remoteName = args
			s.writef("250 %s greets %s", s.srv.Hostname, s.remoteName)

			// RFC 2821 section 4.1.4 specifies that EHLO has the same effect as RSET.
			from = ""
			to = nil
			buffer.Reset()
		case "MAIL":
			match := mailFromRE.FindStringSubmatch(args)
			if match == nil {
				s.writef("501 Syntax error in parameters or arguments (invalid FROM parameter)")
			} else {
				from = match[1]
				s.writef("250 Ok")
			}
			to = nil
			buffer.Reset()
		case "RCPT":
			if from == "" {
				s.writef("503 Bad sequence of commands (MAIL required before RCPT)")
				break
			}

			match := rcptToRE.FindStringSubmatch(args)
			if match == nil {
				s.writef("501 Syntax error in parameters or arguments (invalid TO parameter)")
			} else {
				// RFC 5321 specifies 100 minimum recipients
				if len(to) == 100 {
					s.writef("452 Too many recipients")
				} else {
					to = append(to, match[1])
					s.writef("250 Ok")
				}
			}
		case "DATA":
			if from == "" || to == nil {
				s.writef("503 Bad sequence of commands (MAIL & RCPT required before DATA)")
				break
			}

			s.writef("354 Start mail input; end with <CR><LF>.<CR><LF>")

			// Attempt to read message body from the socket.
			// On error, assume the client has gone away i.e. return from serve().
			data, err := s.readData()
			if err != nil {
				break loop
			}

			// Create Received header & write message body into buffer.
			buffer.Reset()
			buffer.Write(s.makeHeaders(to))
			buffer.Write(data)
			s.writef("250 Ok: queued")

			// Pass mail on to handler.
			if s.srv.Handler != nil {
				go s.srv.Handler(s.conn.RemoteAddr(), from, to, buffer.Bytes())
			}

			// Reset for next mail.
			from = ""
			to = nil
			buffer.Reset()
		case "QUIT":
			s.writef("221 %s %s SMTP Service closing transmission channel", s.srv.Hostname, s.srv.Appname)
			break loop
		case "RSET":
			s.writef("250 Ok")
			from = ""
			to = nil
			buffer.Reset()
		case "NOOP":
			s.writef("250 Ok")
		case "HELP", "VRFY", "EXPN":
			// See RFC 5321 section 4.2.4 for usage of 500 & 502 reply codes
			s.writef("502 Command not implemented")
		default:
			// See RFC 5321 section 4.2.4 for usage of 500 & 502 reply codes
			s.writef("500 Syntax error, command unrecognized")
		}
	}
}

// Wrapper function for writing a complete line to the socket.
func (s *session) writef(format string, args ...interface{}) {
	fmt.Fprintf(s.bw, format+"\r\n", args...)
	s.bw.Flush()
}

// Read a complete line from the socket.
func (s *session) readLine() (string, error) {
	line, err := s.br.ReadString('\n')
	if err != nil {
		return "", err
	}
	line = strings.TrimSpace(line) // Strip trailing \r\n
	return line, err
}

// Parse a line read from the socket.
func (s *session) parseLine(line string) (verb string, args string) {
	if idx := strings.Index(line, " "); idx != -1 {
		verb = strings.ToUpper(line[:idx])
		args = strings.TrimSpace(line[idx+1 : len(line)])
	} else {
		verb = strings.ToUpper(line)
		args = ""
	}
	return verb, args
}

// Read the message data following a DATA command.
func (s *session) readData() ([]byte, error) {
	var data []byte
	for {
		slice, err := s.br.ReadSlice('\n')
		if err != nil {
			return nil, err
		}
		// Handle end of data denoted by lone period (\r\n.\r\n)
		if bytes.Equal(slice, []byte(".\r\n")) {
			break
		}
		// Remove leading period (RFC 5321 section 4.5.2)
		if slice[0] == '.' {
			slice = slice[1:]
		}
		data = append(data, slice...)
	}
	return data, nil
}

// Create the Received header to comply with RFC 2821 section 3.8.2.
// TODO: Work out what to do with multiple to addresses.
func (s *session) makeHeaders(to []string) []byte {
	var buffer bytes.Buffer
	now := time.Now().Format("Mon, _2 Jan 2006 15:04:05 -0700 (MST)")
	buffer.WriteString(fmt.Sprintf("Received: from %s (%s [%s])\r\n", s.remoteName, s.remoteHost, s.remoteIP))
	buffer.WriteString(fmt.Sprintf("        by %s (%s) with SMTP\r\n", s.srv.Hostname, s.srv.Appname))
	buffer.WriteString(fmt.Sprintf("        for <%s>; %s\r\n", to[0], now))
	return buffer.Bytes()
}
<|endoftext|>