<commit_before>package main

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/url"
	"strconv"
	"strings"

	"github.com/agl/xmpp-client/xmpp"
	"golang.org/x/crypto/otr"
	"golang.org/x/crypto/ssh/terminal"
	"golang.org/x/net/proxy"
)

type Config struct {
	filename                      string `json:"-"`
	Account                       string
	Server                        string   `json:",omitempty"`
	Resource                      string   `json:",omitempty"`
	Proxies                       []string `json:",omitempty"`
	Password                      string   `json:",omitempty"`
	Port                          int      `json:",omitempty"`
	PrivateKey                    []byte
	KnownFingerprints             []KnownFingerprint
	RawLogFile                    string   `json:",omitempty"`
	NotifyCommand                 []string `json:",omitempty"`
	IdleSecondsBeforeNotification int      `json:",omitempty"`
	Bell                          bool
	HideStatusUpdates             bool
	UseTor                        bool
	OTRAutoTearDown               bool
	OTRAutoAppendTag              bool
	OTRAutoStartSession           bool
	ServerCertificateSHA256       string   `json:",omitempty"`
	AlwaysEncrypt                 bool     `json:",omitempty"`
	AlwaysEncryptWith             []string `json:",omitempty"`
}

type KnownFingerprint struct {
	UserId         string
	FingerprintHex string
	fingerprint    []byte `json:"-"`
}

func ParseConfig(filename string) (c *Config, err error) {
	contents, err := ioutil.ReadFile(filename)
	if err != nil {
		return
	}

	c = new(Config)
	if err = json.Unmarshal(contents, &c); err != nil {
		return
	}

	c.filename = filename

	for i, known := range c.KnownFingerprints {
		c.KnownFingerprints[i].fingerprint, err = hex.DecodeString(known.FingerprintHex)
		if err != nil {
			err = errors.New("xmpp: failed to parse hex fingerprint for " + known.UserId + ": " + err.Error())
			return
		}
	}

	return
}

func (c *Config) Save() error {
	for i, known := range c.KnownFingerprints {
		c.KnownFingerprints[i].FingerprintHex = hex.EncodeToString(known.fingerprint)
	}

	contents, err := json.MarshalIndent(c, "", "\t")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(c.filename, contents, 0600)
}

func (c *Config) UserIdForFingerprint(fpr []byte) string {
	for _, known := range c.KnownFingerprints {
		if bytes.Equal(fpr, known.fingerprint) {
			return known.UserId
		}
	}

	return ""
}

func (c *Config) HasFingerprint(uid string) bool {
	for _, known := range c.KnownFingerprints {
		if uid == known.UserId {
			return true
		}
	}

	return false
}

func (c *Config) ShouldEncryptTo(uid string) bool {
	if c.AlwaysEncrypt {
		return true
	}

	for _, contact := range c.AlwaysEncryptWith {
		if contact == uid {
			return true
		}
	}
	return false
}

func isYes(s string) bool {
	lower := strings.ToLower(s)
	return lower == "yes" || lower == "y"
}

func enroll(config *Config, term *terminal.Terminal) bool {
	var err error
	warn(term, "Enrolling new config file")

	var domain string
	for {
		term.SetPrompt("Account (i.e. user@example.com, enter to quit): ")
		if config.Account, err = term.ReadLine(); err != nil || len(config.Account) == 0 {
			return false
		}

		parts := strings.SplitN(config.Account, "@", 2)
		if len(parts) != 2 {
			alert(term, "invalid username (want user@domain): "+config.Account)
			continue
		}
		domain = parts[1]
		break
	}

	term.SetPrompt("Resource name (i.e. work, enter for empty): ")
	if config.Resource, err = term.ReadLine(); err != nil {
		return false
	}

	const debugLogFile = "/tmp/xmpp-client-debug.log"
	term.SetPrompt("Enable debug logging to " + debugLogFile + " (y/n)?: ")
	if debugLog, err := term.ReadLine(); err != nil || !isYes(debugLog) {
		info(term, "Not enabling debug logging...")
	} else {
		config.RawLogFile = debugLogFile
		info(term, "Debug logging enabled.")
	}

	term.SetPrompt("Use Tor (y/n)?: ")
	if useTorQuery, err := term.ReadLine(); err != nil || !isYes(useTorQuery) {
		info(term, "Not using Tor...")
		config.UseTor = false
	} else {
		info(term, "Using Tor...")
		config.UseTor = true
	}

	term.SetPrompt("File to import libotr private key from (enter to generate): ")

	var priv otr.PrivateKey
	for {
		importFile, err := term.ReadLine()
		if err != nil {
			return false
		}
		if len(importFile) > 0 {
			privKeyBytes, err := ioutil.ReadFile(importFile)
			if err != nil {
				alert(term, "Failed to open private key file: "+err.Error())
				continue
			}

			if !priv.Import(privKeyBytes) {
				alert(term, "Failed to parse libotr private key file (the parser is pretty simple I'm afraid)")
				continue
			}
			break
		} else {
			info(term, "Generating private key...")
			priv.Generate(rand.Reader)
			break
		}
	}
	config.PrivateKey = priv.Serialize(nil)

	config.OTRAutoAppendTag = true
	config.OTRAutoStartSession = true
	config.OTRAutoTearDown = false

	// List well known Tor hidden services.
	knownTorDomain := map[string]string{
		"jabber.ccc.de":             "okj7xc6j2szr2y75.onion",
		"riseup.net":                "4cjw6cwpeaeppfqz.onion",
		"jabber.calyxinstitute.org": "ijeeynrc6x2uy5ob.onion",
		"jabber.otr.im":             "5rgdtlawqkcplz75.onion",
		"wtfismyip.com":             "ofkztxcohimx34la.onion",
		"rows.io":                   "yz6yiv2hxyagvwy6.onion",
	}

	// Autoconfigure well known Tor hidden services.
	if hiddenService, ok := knownTorDomain[domain]; ok && config.UseTor {
		const torProxyURL = "socks5://127.0.0.1:9050"
		info(term, "It appears that you are using a well known server and we will use its Tor hidden service to connect.")
		config.Server = hiddenService
		config.Port = 5222
		config.Proxies = []string{torProxyURL}
		term.SetPrompt("> ")
		return true
	}

	var proxyStr string
	proxyDefaultPrompt := ", enter for none"
	if config.UseTor {
		proxyDefaultPrompt = ", which is the default"
	}
	term.SetPrompt("Proxy (i.e socks5://127.0.0.1:9050" + proxyDefaultPrompt + "): ")

	for {
		if proxyStr, err = term.ReadLine(); err != nil {
			return false
		}
		if len(proxyStr) == 0 {
			if !config.UseTor {
				break
			} else {
				proxyStr = "socks5://127.0.0.1:9050"
			}
		}
		u, err := url.Parse(proxyStr)
		if err != nil {
			alert(term, "Failed to parse "+proxyStr+" as a URL: "+err.Error())
			continue
		}
		if _, err = proxy.FromURL(u, proxy.Direct); err != nil {
			alert(term, "Failed to parse "+proxyStr+" as a proxy: "+err.Error())
			continue
		}
		break
	}

	if len(proxyStr) > 0 {
		config.Proxies = []string{proxyStr}

		info(term, "Since you selected a proxy, we need to know the server and port to connect to as a SRV lookup would leak information every time.")
		term.SetPrompt("Server (i.e. xmpp.example.com, enter to lookup using unproxied DNS): ")
		if config.Server, err = term.ReadLine(); err != nil {
			return false
		}
		if len(config.Server) == 0 {
			var port uint16
			info(term, "Performing SRV lookup")
			if config.Server, port, err = xmpp.Resolve(domain); err != nil {
				alert(term, "SRV lookup failed: "+err.Error())
				return false
			}
			config.Port = int(port)
			info(term, "Resolved "+config.Server+":"+strconv.Itoa(config.Port))
		} else {
			for {
				term.SetPrompt("Port (enter for 5222): ")
				portStr, err := term.ReadLine()
				if err != nil {
					return false
				}
				if len(portStr) == 0 {
					portStr = "5222"
				}
				if config.Port, err = strconv.Atoi(portStr); err != nil || config.Port <= 0 || config.Port > 65535 {
					info(term, "Port numbers must be 0 < port <= 65535")
					continue
				}
				break
			}
		}
	}

	term.SetPrompt("> ")

	return true
}
<commit_msg>Remove wtfismyip.com due to its shutdown<commit_after>package main

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/url"
	"strconv"
	"strings"

	"github.com/agl/xmpp-client/xmpp"
	"golang.org/x/crypto/otr"
	"golang.org/x/crypto/ssh/terminal"
	"golang.org/x/net/proxy"
)

type Config struct {
	filename                      string `json:"-"`
	Account                       string
	Server                        string   `json:",omitempty"`
	Resource                      string   `json:",omitempty"`
	Proxies                       []string `json:",omitempty"`
	Password                      string   `json:",omitempty"`
	Port                          int      `json:",omitempty"`
	PrivateKey                    []byte
	KnownFingerprints             []KnownFingerprint
	RawLogFile                    string   `json:",omitempty"`
	NotifyCommand                 []string `json:",omitempty"`
	IdleSecondsBeforeNotification int      `json:",omitempty"`
	Bell                          bool
	HideStatusUpdates             bool
	UseTor                        bool
	OTRAutoTearDown               bool
	OTRAutoAppendTag              bool
	OTRAutoStartSession           bool
	ServerCertificateSHA256       string   `json:",omitempty"`
	AlwaysEncrypt                 bool     `json:",omitempty"`
	AlwaysEncryptWith             []string `json:",omitempty"`
}

type KnownFingerprint struct {
	UserId         string
	FingerprintHex string
	fingerprint    []byte `json:"-"`
}

func ParseConfig(filename string) (c *Config, err error) {
	contents, err := ioutil.ReadFile(filename)
	if err != nil {
		return
	}

	c = new(Config)
	if err = json.Unmarshal(contents, &c); err != nil {
		return
	}

	c.filename = filename

	for i, known := range c.KnownFingerprints {
		c.KnownFingerprints[i].fingerprint, err = hex.DecodeString(known.FingerprintHex)
		if err != nil {
			err = errors.New("xmpp: failed to parse hex fingerprint for " + known.UserId + ": " + err.Error())
			return
		}
	}

	return
}

func (c *Config) Save() error {
	for i, known := range c.KnownFingerprints {
		c.KnownFingerprints[i].FingerprintHex = hex.EncodeToString(known.fingerprint)
	}

	contents, err := json.MarshalIndent(c, "", "\t")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(c.filename, contents, 0600)
}

func (c *Config) UserIdForFingerprint(fpr []byte) string {
	for _, known := range c.KnownFingerprints {
		if bytes.Equal(fpr, known.fingerprint) {
			return known.UserId
		}
	}

	return ""
}

func (c *Config) HasFingerprint(uid string) bool {
	for _, known := range c.KnownFingerprints {
		if uid == known.UserId {
			return true
		}
	}

	return false
}

func (c *Config) ShouldEncryptTo(uid string) bool {
	if c.AlwaysEncrypt {
		return true
	}

	for _, contact := range c.AlwaysEncryptWith {
		if contact == uid {
			return true
		}
	}
	return false
}

func isYes(s string) bool {
	lower := strings.ToLower(s)
	return lower == "yes" || lower == "y"
}

func enroll(config *Config, term *terminal.Terminal) bool {
	var err error
	warn(term, "Enrolling new config file")

	var domain string
	for {
		term.SetPrompt("Account (i.e. user@example.com, enter to quit): ")
		if config.Account, err = term.ReadLine(); err != nil || len(config.Account) == 0 {
			return false
		}

		parts := strings.SplitN(config.Account, "@", 2)
		if len(parts) != 2 {
			alert(term, "invalid username (want user@domain): "+config.Account)
			continue
		}
		domain = parts[1]
		break
	}

	term.SetPrompt("Resource name (i.e. work, enter for empty): ")
	if config.Resource, err = term.ReadLine(); err != nil {
		return false
	}

	const debugLogFile = "/tmp/xmpp-client-debug.log"
	term.SetPrompt("Enable debug logging to " + debugLogFile + " (y/n)?: ")
	if debugLog, err := term.ReadLine(); err != nil || !isYes(debugLog) {
		info(term, "Not enabling debug logging...")
	} else {
		config.RawLogFile = debugLogFile
		info(term, "Debug logging enabled.")
	}

	term.SetPrompt("Use Tor (y/n)?: ")
	if useTorQuery, err := term.ReadLine(); err != nil || !isYes(useTorQuery) {
		info(term, "Not using Tor...")
		config.UseTor = false
	} else {
		info(term, "Using Tor...")
		config.UseTor = true
	}

	term.SetPrompt("File to import libotr private key from (enter to generate): ")

	var priv otr.PrivateKey
	for {
		importFile, err := term.ReadLine()
		if err != nil {
			return false
		}
		if len(importFile) > 0 {
			privKeyBytes, err := ioutil.ReadFile(importFile)
			if err != nil {
				alert(term, "Failed to open private key file: "+err.Error())
				continue
			}

			if !priv.Import(privKeyBytes) {
				alert(term, "Failed to parse libotr private key file (the parser is pretty simple I'm afraid)")
				continue
			}
			break
		} else {
			info(term, "Generating private key...")
			priv.Generate(rand.Reader)
			break
		}
	}
	config.PrivateKey = priv.Serialize(nil)

	config.OTRAutoAppendTag = true
	config.OTRAutoStartSession = true
	config.OTRAutoTearDown = false

	// List well known Tor hidden services.
	knownTorDomain := map[string]string{
		"jabber.ccc.de":             "okj7xc6j2szr2y75.onion",
		"riseup.net":                "4cjw6cwpeaeppfqz.onion",
		"jabber.calyxinstitute.org": "ijeeynrc6x2uy5ob.onion",
		"jabber.otr.im":             "5rgdtlawqkcplz75.onion",
		"rows.io":                   "yz6yiv2hxyagvwy6.onion",
	}

	// Autoconfigure well known Tor hidden services.
	if hiddenService, ok := knownTorDomain[domain]; ok && config.UseTor {
		const torProxyURL = "socks5://127.0.0.1:9050"
		info(term, "It appears that you are using a well known server and we will use its Tor hidden service to connect.")
		config.Server = hiddenService
		config.Port = 5222
		config.Proxies = []string{torProxyURL}
		term.SetPrompt("> ")
		return true
	}

	var proxyStr string
	proxyDefaultPrompt := ", enter for none"
	if config.UseTor {
		proxyDefaultPrompt = ", which is the default"
	}
	term.SetPrompt("Proxy (i.e socks5://127.0.0.1:9050" + proxyDefaultPrompt + "): ")

	for {
		if proxyStr, err = term.ReadLine(); err != nil {
			return false
		}
		if len(proxyStr) == 0 {
			if !config.UseTor {
				break
			} else {
				proxyStr = "socks5://127.0.0.1:9050"
			}
		}
		u, err := url.Parse(proxyStr)
		if err != nil {
			alert(term, "Failed to parse "+proxyStr+" as a URL: "+err.Error())
			continue
		}
		if _, err = proxy.FromURL(u, proxy.Direct); err != nil {
			alert(term, "Failed to parse "+proxyStr+" as a proxy: "+err.Error())
			continue
		}
		break
	}

	if len(proxyStr) > 0 {
		config.Proxies = []string{proxyStr}

		info(term, "Since you selected a proxy, we need to know the server and port to connect to as a SRV lookup would leak information every time.")
		term.SetPrompt("Server (i.e. xmpp.example.com, enter to lookup using unproxied DNS): ")
		if config.Server, err = term.ReadLine(); err != nil {
			return false
		}
		if len(config.Server) == 0 {
			var port uint16
			info(term, "Performing SRV lookup")
			if config.Server, port, err = xmpp.Resolve(domain); err != nil {
				alert(term, "SRV lookup failed: "+err.Error())
				return false
			}
			config.Port = int(port)
			info(term, "Resolved "+config.Server+":"+strconv.Itoa(config.Port))
		} else {
			for {
				term.SetPrompt("Port (enter for 5222): ")
				portStr, err := term.ReadLine()
				if err != nil {
					return false
				}
				if len(portStr) == 0 {
					portStr = "5222"
				}
				if config.Port, err = strconv.Atoi(portStr); err != nil || config.Port <= 0 || config.Port > 65535 {
					info(term, "Port numbers must be 0 < port <= 65535")
					continue
				}
				break
			}
		}
	}

	term.SetPrompt("> ")

	return true
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"launchpad.net/goyaml"
)

type Config struct {
	APIEndpoint   string   `yaml:"api_endpoint"`
	APIKey        string   `yaml:"api_key"`
	ProjectBranch string   `yaml:"project_branch"`
	IgnoredPaths  []string `yaml:"ignored_paths"`
}

// Load config from config file
//
func NewConfig(config_data []byte) (*Config, error) {
	config := &Config{
		APIEndpoint:   "https://gemnasium.com/api/v3",
		ProjectBranch: "master",
	}
	goyaml.Unmarshal(config_data, config)
	return config, nil

}

func LoadConfigFile(filepath string) (*Config, error) {
	config_data, err := ioutil.ReadFile(filepath)
	if err != nil {
		return nil, err
	}
	return NewConfig(config_data)
}
<commit_msg>goyaml has moved to github<commit_after>package main

import (
	"io/ioutil"

	"gopkg.in/yaml.v1"
)

type Config struct {
	APIEndpoint   string   `yaml:"api_endpoint"`
	APIKey        string   `yaml:"api_key"`
	ProjectBranch string   `yaml:"project_branch"`
	IgnoredPaths  []string `yaml:"ignored_paths"`
}

// Load config from config file
//
func NewConfig(config_data []byte) (*Config, error) {
	config := &Config{
		APIEndpoint:   "https://gemnasium.com/api/v3",
		ProjectBranch: "master",
	}
	yaml.Unmarshal(config_data, config)
	return config, nil

}

func LoadConfigFile(filepath string) (*Config, error) {
	config_data, err := ioutil.ReadFile(filepath)
	if err != nil {
		return nil, err
	}
	return NewConfig(config_data)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"github.com/miekg/dns"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path"
	"runtime"
	"runtime/debug"
	"sort"
	"strconv"
	"strings"
	"time"
)

var configLastRead = map[string]time.Time{}

func configReader(dirName string, Zones Zones) {
	for {
		configReadDir(dirName, Zones)
		time.Sleep(5 * time.Second)
	}
}

func configReadDir(dirName string, Zones Zones) {
	dir, err := ioutil.ReadDir(dirName)
	if err != nil {
		panic(err)
	}

	seenFiles := map[string]bool{}

	for _, file := range dir {
		fileName := file.Name()
		if !strings.HasSuffix(strings.ToLower(fileName), ".json") {
			continue
		}

		seenFiles[fileName] = true

		if lastRead, ok := configLastRead[fileName]; !ok || file.ModTime().After(lastRead) {
			log.Println("Updated file, going to read", fileName)
			configLastRead[fileName] = file.ModTime()
			zoneName := fileName[0:strings.LastIndex(fileName, ".")]
			//log.Println("FILE:", i, file, zoneName)
			runtime.GC()
			config, err := readZoneFile(zoneName, path.Join(dirName, fileName))
			if config == nil || err != nil {
				log.Println("error reading file: ", err)
			}
			if config != nil && err == nil {
				Zones[zoneName] = config
				dns.HandleFunc(zoneName, setupServerFunc(config))
				runtime.GC()
			}
		}

		// TODO(ask) Disable zones not seen in two subsequent runs
	}
}

func setupPgeodnsZone(Zones Zones) {
	zoneName := "pgeodns"
	Zone := new(Zone)
	Zone.Labels = make(labels)
	Zone.Origin = zoneName
	Zone.LenLabels = dns.LenLabels(Zone.Origin)
	label := new(Label)
	label.Records = make(map[uint16]Records)
	label.Weight = make(map[uint16]int)
	Zone.Labels[""] = label
	setupSOA(Zone)
	Zones[zoneName] = Zone
	dns.HandleFunc(zoneName, setupServerFunc(Zone))
}

func readZoneFile(zoneName, fileName string) (*Zone, error) {
	defer func() {
		if err := recover(); err != nil {
			log.Printf("reading %s failed: %s", zoneName, err)
			debug.PrintStack()
		}
	}()

	fh, err := os.Open(fileName)
	if err != nil {
		log.Println("Could not read ", fileName, ": ", err)
		panic(err)
	}

	Zone := new(Zone)
	Zone.Labels = make(labels)
	Zone.Origin = zoneName
	Zone.LenLabels = dns.LenLabels(Zone.Origin)
	Zone.Options.Ttl = 120
	Zone.Options.MaxHosts = 2

	if err == nil {
		var objmap map[string]interface{}
		decoder := json.NewDecoder(fh)
		err := decoder.Decode(&objmap)
		if err != nil {
			panic(err)
		}
		//log.Println(objmap)

		var data map[string]interface{}

		for k, v := range objmap {
			//log.Printf("k: %s v: %#v, T: %T\n", k, v, v)

			switch k {
			case "ttl", "serial", "max_hosts":
				switch option := k; option {
				case "ttl":
					Zone.Options.Ttl = valueToInt(v)
				case "serial":
					Zone.Options.Serial = valueToInt(v)
				case "max_hosts":
					Zone.Options.MaxHosts = valueToInt(v)
				}
				continue

			case "data":
				data = v.(map[string]interface{})
			}
		}

		setupZoneData(data, Zone)

	}

	//log.Printf("ZO T: %T %s\n", Zones["0.us"], Zones["0.us"])

	//log.Println("IP", string(Zone.Regions["0.us"].IPv4[0].ip))

	return Zone, nil
}

func setupZoneData(data map[string]interface{}, Zone *Zone) {

	recordTypes := map[string]uint16{
		"a":     dns.TypeA,
		"aaaa":  dns.TypeAAAA,
		"ns":    dns.TypeNS,
		"cname": dns.TypeCNAME,
		"mx":    dns.TypeMX,
		"alias": dns.TypeMF,
	}

	for dk, dv_inter := range data {

		dv := dv_inter.(map[string]interface{})

		//log.Printf("K %s V %s TYPE-V %T\n", dk, dv, dv)

		dk = strings.ToLower(dk)
		Zone.Labels[dk] = new(Label)
		label := Zone.Labels[dk]
		label.Label = dk
		label.Ttl = Zone.Options.Ttl
		label.MaxHosts = Zone.Options.MaxHosts

		if ttl, ok := dv["ttl"]; ok {
			label.Ttl = valueToInt(ttl)
		}

		if maxHosts, ok := dv["max_hosts"]; ok {
			label.MaxHosts = valueToInt(maxHosts)
		}

		for rType, dnsType := range recordTypes {

			rdata := dv[rType]

			if rdata == nil {
				//log.Printf("No %s records for label %s\n", rType, dk)
				continue
			}

			//log.Printf("rdata %s TYPE-R %T\n", rdata, rdata)

			records := make(map[string][]interface{})

			switch rdata.(type) {
			case map[string]interface{}:
				// Handle NS map syntax, map[ns2.example.net:<nil> ns1.example.net:<nil>]
				tmp := make([]interface{}, 0)
				for rdata_k, rdata_v := range rdata.(map[string]interface{}) {
					if rdata_v == nil {
						rdata_v = ""
					}
					tmp = append(tmp, []string{rdata_k, rdata_v.(string)})
				}
				records[rType] = tmp
			case string:
				// CNAME and alias
				tmp := make([]interface{}, 1)
				tmp[0] = rdata.(string)
				records[rType] = tmp
			default:
				records[rType] = rdata.([]interface{})
			}

			//log.Printf("RECORDS %s TYPE-REC %T\n", Records, Records)

			if label.Records == nil {
				label.Records = make(map[uint16]Records)
				label.Weight = make(map[uint16]int)
			}

			label.Records[dnsType] = make(Records, len(records[rType]))

			for i := 0; i < len(records[rType]); i++ {

				//log.Printf("RT %T %#v\n", records[rType][i], records[rType][i])

				record := new(Record)

				var h dns.RR_Header
				// log.Println("TTL OPTIONS", Zone.Options.Ttl)
				h.Ttl = uint32(label.Ttl)
				h.Class = dns.ClassINET
				h.Rrtype = dnsType
				h.Name = label.Label + "." + Zone.Origin + "."

				switch dnsType {
				case dns.TypeA, dns.TypeAAAA:
					rec := records[rType][i].([]interface{})
					ip := rec[0].(string)
					var err error
					switch rec[1].(type) {
					case string:
						record.Weight, err = strconv.Atoi(rec[1].(string))
						if err != nil {
							panic("Error converting weight to integer")
						}
					case float64:
						record.Weight = int(rec[1].(float64))
					}
					switch dnsType {
					case dns.TypeA:
						if x := net.ParseIP(ip); x != nil {
							record.RR = &dns.RR_A{Hdr: h, A: x}
							break
						}
						panic("Bad A record")
					case dns.TypeAAAA:
						if x := net.ParseIP(ip); x != nil {
							record.RR = &dns.RR_AAAA{Hdr: h, AAAA: x}
							break
						}
						panic("Bad AAAA record")
					}

				case dns.TypeMX:
					rec := records[rType][i].(map[string]interface{})
					pref := uint16(0)
					mx := rec["mx"].(string)
					if !strings.HasSuffix(mx, ".") {
						mx = mx + "."
					}
					if rec["weight"] != nil {
						record.Weight = valueToInt(rec["weight"])
					}
					if rec["preference"] != nil {
						pref = uint16(valueToInt(rec["preference"]))
					}
					record.RR = &dns.RR_MX{
						Hdr:  h,
						Mx:   mx,
						Pref: pref}

				case dns.TypeCNAME:
					rec := records[rType][i]
					record.RR = &dns.RR_CNAME{Hdr: h, Target: dns.Fqdn(rec.(string))}

				case dns.TypeMF:
					rec := records[rType][i]
					// MF records (how we store aliases) are not FQDNs
					record.RR = &dns.RR_MF{Hdr: h, Mf: rec.(string)}

				case dns.TypeNS:
					rec := records[rType][i]
					if h.Ttl < 86400 {
						h.Ttl = 86400
					}

					var ns string

					switch rec.(type) {
					case string:
						ns = rec.(string)
					case []string:
						recl := rec.([]string)
						ns = recl[0]
						if len(recl[1]) > 0 {
							log.Println("NS records with names syntax not supported")
						}
					default:
						log.Printf("Data: %T %#v\n", rec, rec)
						panic("Unrecognized NS format/syntax")
					}

					rr := &dns.RR_NS{Hdr: h, Ns: dns.Fqdn(ns)}

					record.RR = rr

				default:
					log.Println("type:", rType)
					panic("Don't know how to handle this type")
				}

				if record.RR == nil {
					panic("record.RR is nil")
				}

				label.Weight[dnsType] += record.Weight
				label.Records[dnsType][i] = *record
			}
			if label.Weight[dnsType] > 0 {
				sort.Sort(RecordsByWeight{label.Records[dnsType]})
			}
		}
	}

	setupSOA(Zone)

	//log.Println(Zones[k])
}

func setupSOA(Zone *Zone) {
	label := Zone.Labels[""]

	primaryNs := "ns"

	if record, ok := label.Records[dns.TypeNS]; ok {
		primaryNs = record[0].RR.(*dns.RR_NS).Ns
	}

	s := Zone.Origin + ". 3600 IN SOA " +
		primaryNs + " support.bitnames.com. " +
		strconv.Itoa(Zone.Options.Serial) +
		" 5400 5400 2419200 " +
		strconv.Itoa(Zone.Options.Ttl)

	log.Println("SOA: ", s)

	rr, err := dns.NewRR(s)

	if err != nil {
		log.Println("SOA Error", err)
		panic("Could not setup SOA")
	}

	record := Record{RR: rr}

	label.Records[dns.TypeSOA] = make([]Record, 1)
	label.Records[dns.TypeSOA][0] = record

}

func valueToInt(v interface{}) (rv int) {
	switch v.(type) {
	case string:
		i, err := strconv.Atoi(v.(string))
		if err != nil {
			panic("Error converting weight to integer")
		}
		rv = i
	case float64:
		rv = int(v.(float64))
	default:
		log.Println("Can't convert", v, "to integer")
		panic("Can't convert value")
	}
	return rv
}
<commit_msg>Update to work with latest dns library<commit_after>package main

import (
	"encoding/json"
	"github.com/miekg/dns"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path"
	"runtime"
	"runtime/debug"
	"sort"
	"strconv"
	"strings"
	"time"
)

var configLastRead = map[string]time.Time{}

func configReader(dirName string, Zones Zones) {
	for {
		configReadDir(dirName, Zones)
		time.Sleep(5 * time.Second)
	}
}

func configReadDir(dirName string, Zones Zones) {
	dir, err := ioutil.ReadDir(dirName)
	if err != nil {
		panic(err)
	}

	seenFiles := map[string]bool{}

	for _, file := range dir {
		fileName := file.Name()
		if !strings.HasSuffix(strings.ToLower(fileName), ".json") {
			continue
		}

		seenFiles[fileName] = true

		if lastRead, ok := configLastRead[fileName]; !ok || file.ModTime().After(lastRead) {
			log.Println("Updated file, going to read", fileName)
			configLastRead[fileName] = file.ModTime()
			zoneName := fileName[0:strings.LastIndex(fileName, ".")]
			//log.Println("FILE:", i, file, zoneName)
			runtime.GC()
			config, err := readZoneFile(zoneName, path.Join(dirName, fileName))
			if config == nil || err != nil {
				log.Println("error reading file: ", err)
			}
			if config != nil && err == nil {
				Zones[zoneName] = config
				dns.HandleFunc(zoneName, setupServerFunc(config))
				runtime.GC()
			}
		}

		// TODO(ask) Disable zones not seen in two subsequent runs
	}
}

func setupPgeodnsZone(Zones Zones) {
	zoneName := "pgeodns"
	Zone := new(Zone)
	Zone.Labels = make(labels)
	Zone.Origin = zoneName
	Zone.LenLabels = dns.LenLabels(Zone.Origin)
	label := new(Label)
	label.Records = make(map[uint16]Records)
	label.Weight = make(map[uint16]int)
	Zone.Labels[""] = label
	setupSOA(Zone)
	Zones[zoneName] = Zone
	dns.HandleFunc(zoneName, setupServerFunc(Zone))
}

func readZoneFile(zoneName, fileName string) (*Zone, error) {
	defer func() {
		if err := recover(); err != nil {
			log.Printf("reading %s failed: %s", zoneName, err)
			debug.PrintStack()
		}
	}()

	fh, err := os.Open(fileName)
	if err != nil {
		log.Println("Could not read ", fileName, ": ", err)
		panic(err)
	}

	Zone := new(Zone)
	Zone.Labels = make(labels)
	Zone.Origin = zoneName
	Zone.LenLabels = dns.LenLabels(Zone.Origin)
	Zone.Options.Ttl = 120
	Zone.Options.MaxHosts = 2

	if err == nil {
		var objmap map[string]interface{}
		decoder := json.NewDecoder(fh)
		err := decoder.Decode(&objmap)
		if err != nil {
			panic(err)
		}
		//log.Println(objmap)

		var data map[string]interface{}

		for k, v := range objmap {
			//log.Printf("k: %s v: %#v, T: %T\n", k, v, v)

			switch k {
			case "ttl", "serial", "max_hosts":
				switch option := k; option {
				case "ttl":
					Zone.Options.Ttl = valueToInt(v)
				case "serial":
					Zone.Options.Serial = valueToInt(v)
				case "max_hosts":
					Zone.Options.MaxHosts = valueToInt(v)
				}
				continue

			case "data":
				data = v.(map[string]interface{})
			}
		}

		setupZoneData(data, Zone)

	}

	//log.Printf("ZO T: %T %s\n", Zones["0.us"], Zones["0.us"])

	//log.Println("IP", string(Zone.Regions["0.us"].IPv4[0].ip))

	return Zone, nil
}

func setupZoneData(data map[string]interface{}, Zone *Zone) {

	recordTypes := map[string]uint16{
		"a":     dns.TypeA,
		"aaaa":  dns.TypeAAAA,
		"ns":    dns.TypeNS,
		"cname": dns.TypeCNAME,
		"mx":    dns.TypeMX,
		"alias": dns.TypeMF,
	}

	for dk, dv_inter := range data {

		dv := dv_inter.(map[string]interface{})

		//log.Printf("K %s V %s TYPE-V %T\n", dk, dv, dv)

		dk = strings.ToLower(dk)
		Zone.Labels[dk] = new(Label)
		label := Zone.Labels[dk]
		label.Label = dk
		label.Ttl = Zone.Options.Ttl
		label.MaxHosts = Zone.Options.MaxHosts

		if ttl, ok := dv["ttl"]; ok {
			label.Ttl = valueToInt(ttl)
		}

		if maxHosts, ok := dv["max_hosts"]; ok {
			label.MaxHosts = valueToInt(maxHosts)
		}

		for rType, dnsType := range recordTypes {

			rdata := dv[rType]

			if rdata == nil {
				//log.Printf("No %s records for label %s\n", rType, dk)
				continue
			}

			//log.Printf("rdata %s TYPE-R %T\n", rdata, rdata)

			records := make(map[string][]interface{})

			switch rdata.(type) {
			case map[string]interface{}:
				// Handle NS map syntax, map[ns2.example.net:<nil> ns1.example.net:<nil>]
				tmp := make([]interface{}, 0)
				for rdata_k, rdata_v := range rdata.(map[string]interface{}) {
					if rdata_v == nil {
						rdata_v = ""
					}
					tmp = append(tmp, []string{rdata_k, rdata_v.(string)})
				}
				records[rType] = tmp
			case string:
				// CNAME and alias
				tmp := make([]interface{}, 1)
				tmp[0] = rdata.(string)
				records[rType] = tmp
			default:
				records[rType] = rdata.([]interface{})
			}

			//log.Printf("RECORDS %s TYPE-REC %T\n", Records, Records)

			if label.Records == nil {
				label.Records = make(map[uint16]Records)
				label.Weight = make(map[uint16]int)
			}

			label.Records[dnsType] = make(Records, len(records[rType]))

			for i := 0; i < len(records[rType]); i++ {

				//log.Printf("RT %T %#v\n", records[rType][i], records[rType][i])

				record := new(Record)

				var h dns.RR_Header
				// log.Println("TTL OPTIONS", Zone.Options.Ttl)
				h.Ttl = uint32(label.Ttl)
				h.Class = dns.ClassINET
				h.Rrtype = dnsType
				h.Name = label.Label + "." + Zone.Origin + "."

				switch dnsType {
				case dns.TypeA, dns.TypeAAAA:
					rec := records[rType][i].([]interface{})
					ip := rec[0].(string)
					var err error
					switch rec[1].(type) {
					case string:
						record.Weight, err = strconv.Atoi(rec[1].(string))
						if err != nil {
							panic("Error converting weight to integer")
						}
					case float64:
						record.Weight = int(rec[1].(float64))
					}
					switch dnsType {
					case dns.TypeA:
						if x := net.ParseIP(ip); x != nil {
							record.RR = &dns.RR_A{Hdr: h, A: x}
							break
						}
						panic("Bad A record")
					case dns.TypeAAAA:
						if x := net.ParseIP(ip); x != nil {
							record.RR = &dns.RR_AAAA{Hdr: h, AAAA: x}
							break
						}
						panic("Bad AAAA record")
					}

				case dns.TypeMX:
					rec := records[rType][i].(map[string]interface{})
					pref := uint16(0)
					mx := rec["mx"].(string)
					if !strings.HasSuffix(mx, ".") {
						mx = mx + "."
					}
					if rec["weight"] != nil {
						record.Weight = valueToInt(rec["weight"])
					}
					if rec["preference"] != nil {
						pref = uint16(valueToInt(rec["preference"]))
					}
					record.RR = &dns.RR_MX{
						Hdr:        h,
						Mx:         mx,
						Preference: pref}

				case dns.TypeCNAME:
					rec := records[rType][i]
					record.RR = &dns.RR_CNAME{Hdr: h, Target: dns.Fqdn(rec.(string))}

				case dns.TypeMF:
					rec := records[rType][i]
					// MF records (how we store aliases) are not FQDNs
					record.RR = &dns.RR_MF{Hdr: h, Mf: rec.(string)}

				case dns.TypeNS:
					rec := records[rType][i]
					if h.Ttl < 86400 {
						h.Ttl = 86400
					}

					var ns string

					switch rec.(type) {
					case string:
						ns = rec.(string)
					case []string:
						recl := rec.([]string)
						ns = recl[0]
						if len(recl[1]) > 0 {
							log.Println("NS records with names syntax not supported")
						}
					default:
						log.Printf("Data: %T %#v\n", rec, rec)
						panic("Unrecognized NS format/syntax")
					}

					rr := &dns.RR_NS{Hdr: h, Ns: dns.Fqdn(ns)}

					record.RR = rr

				default:
					log.Println("type:", rType)
					panic("Don't know how to handle this type")
				}

				if record.RR == nil {
					panic("record.RR is nil")
				}

				label.Weight[dnsType] += record.Weight
				label.Records[dnsType][i] = *record
			}
			if label.Weight[dnsType] > 0 {
				sort.Sort(RecordsByWeight{label.Records[dnsType]})
			}
		}
	}

	setupSOA(Zone)

	//log.Println(Zones[k])
}

func setupSOA(Zone *Zone) {
	label := Zone.Labels[""]

	primaryNs := "ns"

	if record, ok := label.Records[dns.TypeNS]; ok {
		primaryNs = record[0].RR.(*dns.RR_NS).Ns
	}

	s := Zone.Origin + ". 3600 IN SOA " +
		primaryNs + " support.bitnames.com. " +
		strconv.Itoa(Zone.Options.Serial) +
		" 5400 5400 2419200 " +
		strconv.Itoa(Zone.Options.Ttl)

	log.Println("SOA: ", s)

	rr, err := dns.NewRR(s)

	if err != nil {
		log.Println("SOA Error", err)
		panic("Could not setup SOA")
	}

	record := Record{RR: rr}

	label.Records[dns.TypeSOA] = make([]Record, 1)
	label.Records[dns.TypeSOA][0] = record

}

func valueToInt(v interface{}) (rv int) {
	switch v.(type) {
	case string:
		i, err := strconv.Atoi(v.(string))
		if err != nil {
			panic("Error converting weight to integer")
		}
		rv = i
	case float64:
		rv = int(v.(float64))
	default:
		log.Println("Can't convert", v, "to integer")
		panic("Can't convert value")
	}
	return rv
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/pion/rtcp"
	"github.com/pion/webrtc/v2"
)

type TrackEventType uint32

const (
	TrackEventTypeAdd = iota + 1
	TrackEventTypeRemove
)

type TrackEvent struct {
	ClientID string
	Track    *webrtc.Track
	Type     TrackEventType
}

type TrackMetadata struct {
	Mid      string `json:"mid"`
	UserID   string `json:"userId"`
	StreamID string `json:"streamId"`
	Kind     string `json:"kind"`
}

type TrackInfo struct {
	RTPTransceiver *webrtc.RTPTransceiver
	RTPSender      *webrtc.RTPSender
	TrackMetadata  TrackMetadata
}

type trackListener struct {
	log              Logger
	clientID         string
	peerConnection   *webrtc.PeerConnection
	localTracks      []*webrtc.Track
	localTracksMu    sync.RWMutex
	trackInfoByTrack map[*webrtc.Track]TrackInfo
	onTrackEvent     func(TrackEvent)
	mu               sync.RWMutex
	pliInterval      time.Duration
	ssrcMaxBitrates  map[uint32]uint64
}

func newTrackListener(
	loggerFactory LoggerFactory,
	clientID string,
	peerConnection *webrtc.PeerConnection,
	onTrackEvent func(TrackEvent),
) *trackListener {
	p := &trackListener{
		log:              loggerFactory.GetLogger("tracklistener"),
		clientID:         clientID,
		peerConnection:   peerConnection,
		trackInfoByTrack: map[*webrtc.Track]TrackInfo{},
		onTrackEvent:     onTrackEvent,
		ssrcMaxBitrates:  map[uint32]uint64{},
	}

	p.log.Printf("[%s] Setting PeerConnection.OnTrack listener", clientID)
	peerConnection.OnTrack(p.handleTrack)

	return p
}

func (p *trackListener) ClientID() string {
	return p.clientID
}

// GetTracksMetadata gets metadata of the sending tracks with updated Mid
func (p *trackListener) GetTracksMetadata() (metadata []TrackMetadata) {
	p.mu.RLock()
	defer p.mu.RUnlock()

	metadata = make([]TrackMetadata, 0)

	for _, trackInfo := range p.trackInfoByTrack {
		m := trackInfo.TrackMetadata
		m.Mid = trackInfo.RTPTransceiver.Mid()
		metadata = append(metadata, m)
	}
	return
}

func (p *trackListener) WriteRTCP(anyPacket rtcp.Packet) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	switch pkt := anyPacket.(type) {
	case *rtcp.PictureLossIndication:
		prometheusRTCPPacketsSent.Inc()
		return p.peerConnection.WriteRTCP([]rtcp.Packet{pkt})
	case *rtcp.SourceDescription:
	case *rtcp.ReceiverEstimatedMaximumBitrate:
		bitrate, ok := p.ssrcMaxBitrates[pkt.SenderSSRC]
		if !ok || pkt.Bitrate < bitrate {
			p.ssrcMaxBitrates[pkt.SenderSSRC] = bitrate
			prometheusRTCPPacketsSent.Inc()
			return p.peerConnection.WriteRTCP([]rtcp.Packet{pkt})
		}
	case *rtcp.ReceiverReport:
	case *rtcp.SenderReport:
	default:
		p.log.Printf("[%s] Got unhandled RTCP pkt for track: %d (%T)", p.clientID, pkt.DestinationSSRC(), pkt)
	}

	return nil
}

func (p *trackListener) AddTrack(sourceClientID string, track *webrtc.Track) (chan rtcp.Packet, error) {
	p.localTracksMu.Lock()
	defer p.localTracksMu.Unlock()

	p.log.Printf("[%s] peer.AddTrack: %d", p.clientID, track.SSRC())
	rtpSender, err := p.peerConnection.AddTrack(track)

	var transceiver *webrtc.RTPTransceiver
	for _, tr := range p.peerConnection.GetTransceivers() {
		if tr.Sender() == rtpSender {
			transceiver = tr
			break
		}
	}

	rtcpCh := make(chan rtcp.Packet)

	if err != nil {
		close(rtcpCh)
		return rtcpCh, fmt.Errorf("[%s] peer.AddTrack: error adding track: %d: %s", p.clientID, track.SSRC(), err)
	}

	go func() {
		for {
			rtcps, err := rtpSender.ReadRTCP()
			if err != nil {
				p.log.Printf("[%s] RTCP stream for sender track: %d has ended: %s",
					p.clientID,
					track.SSRC(),
					err,
				)
				close(rtcpCh)
				break
			}
			for _, pkt := range rtcps {
				prometheusRTCPPacketsReceived.Inc()
				rtcpCh <- pkt
			}
		}
	}()

	p.trackInfoByTrack[track] = TrackInfo{
		RTPSender:      rtpSender,
		RTPTransceiver: transceiver,
		TrackMetadata: TrackMetadata{
			Mid:      "",
			Kind:     track.Kind().String(),
			UserID:   sourceClientID,
			StreamID: track.Label(),
		},
	}
	return rtcpCh, nil
}

func (p *trackListener) RemoveTrack(track *webrtc.Track) error {
	p.localTracksMu.Lock()
	defer p.localTracksMu.Unlock()
	p.log.Printf("[%s] peer.RemoveTrack: %d", p.clientID, track.SSRC())
	trackInfo, ok := p.trackInfoByTrack[track]
	if !ok {
		return fmt.Errorf("[%s] peer.RemoveTrack: cannot find sender for track: %d", p.clientID, track.SSRC())
	}
	delete(p.trackInfoByTrack, track)
	delete(p.ssrcMaxBitrates, track.SSRC())
	return p.peerConnection.RemoveTrack(trackInfo.RTPSender)
}

func (p *trackListener) handleTrack(remoteTrack *webrtc.Track, receiver *webrtc.RTPReceiver) {
	p.log.Printf("[%s] peer.handleTrack (id: %s, label: %s, type: %s, ssrc: %d)",
		p.clientID, remoteTrack.ID(), remoteTrack.Label(), remoteTrack.Kind(), remoteTrack.SSRC())
	localTrack, err := p.startCopyingTrack(remoteTrack, receiver)
	if err != nil {
		p.log.Printf("Error copying remote track: %s", err)
		return
	}
	p.localTracksMu.Lock()
	p.localTracks = append(p.localTracks, localTrack)
	p.localTracksMu.Unlock()

	p.log.Printf("[%s] peer.handleTrack add track to list of local tracks: %d", p.clientID, localTrack.SSRC())

	p.sendTrackEvent(TrackEvent{p.clientID, localTrack, TrackEventTypeAdd})
}

func (p *trackListener) sendTrackEvent(t TrackEvent) {
	go p.onTrackEvent(t)
}

func (p *trackListener) Tracks() []*webrtc.Track {
	return p.localTracks
}

func (p *trackListener) startCopyingTrack(remoteTrack *webrtc.Track, receiver *webrtc.RTPReceiver) (*webrtc.Track, error) {
	remoteTrackID := remoteTrack.ID()
	if remoteTrackID == "" {
		remoteTrackID = NewUUIDBase62()
	}
	// this is the media stream ID we add the p.clientID in the string to know
	// which user the video came from and the remoteTrack.Label() so we can
	// associate audio/video tracks from the same MediaStream
	remoteTrackLabel := remoteTrack.Label()
	if remoteTrackLabel == "" {
		remoteTrackLabel = NewUUIDBase62()
	}
	localTrackLabel := "sfu_" + p.clientID + "_" + remoteTrackLabel

	localTrackID := "sfu_" + remoteTrackID
	p.log.Printf("[%s] peer.startCopyingTrack: %d", p.clientID, remoteTrack.SSRC())

	ssrc := remoteTrack.SSRC()
	// Create a local track, all our SFU clients will be fed via this track
	localTrack, err := p.peerConnection.NewTrack(remoteTrack.PayloadType(), ssrc, localTrackID, localTrackLabel)
	if err != nil {
		err = fmt.Errorf("[%s] peer.startCopyingTrack: error creating new track: %d, error: %s", p.clientID, remoteTrack.SSRC(), err)
		return nil, err
	}

	// Send a PLI on an interval so that the publisher is pushing a keyframe every rtcpPLIInterval
	// This can be less wasteful by processing incoming RTCP events, then we would emit a NACK/PLI when a viewer requests it

	var ticker *time.Ticker
	if p.pliInterval > 0 {
		ticker = time.NewTicker(p.pliInterval)
		go func() {
			writeRTCP := func() {
				err := p.peerConnection.WriteRTCP(
					[]rtcp.Packet{
						&rtcp.PictureLossIndication{
							MediaSSRC: ssrc,
						},
					},
				)
				if err != nil {
					p.log.Printf("[%s] Error sending rtcp PLI for local track: %d: %s",
						p.clientID,
						localTrack.SSRC(),
						err,
					)
				}
			}
			for range ticker.C {
				writeRTCP()
			}
		}()
	}

	go func() {
		start := time.Now()
		prometheusWebRTCTracksTotal.Inc()
		prometheusWebRTCTracksActive.Inc()
		defer func() {
			prometheusWebRTCTracksActive.Dec()
			prometheusWebRTCTracksDuration.Observe(time.Now().Sub(start).Seconds())
		}()
		defer p.sendTrackEvent(TrackEvent{p.clientID, localTrack, TrackEventTypeRemove})
		if ticker != nil {
			defer ticker.Stop()
		}
		for {
			pkt, err := remoteTrack.ReadRTP()
			if err != nil {
				p.log.Printf(
					"[%s] Remote track has ended: %d: %s",
					p.clientID,
					remoteTrack.SSRC(),
					err,
				)
				return
			}

			prometheusRTPPacketsReceived.Inc()

			// ErrClosedPipe means we don't have any subscribers, this is ok if no peers have connected yet
			err = localTrack.WriteRTP(pkt)
			if err != nil && err != io.ErrClosedPipe {
				p.log.Printf(
					"[%s] Error writing to local track: %d: %s",
					p.clientID,
					localTrack.SSRC(),
					err,
				)
				return
			}
		}
	}()

	return localTrack, nil
}
<commit_msg>Mention RTP in log entry about ending track<commit_after>package server

import (
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/pion/rtcp"
	"github.com/pion/webrtc/v2"
)

type TrackEventType uint32

const (
	TrackEventTypeAdd = iota + 1
	TrackEventTypeRemove
)

type TrackEvent struct {
	ClientID string
	Track    *webrtc.Track
	Type     TrackEventType
}

type TrackMetadata struct {
	Mid      string `json:"mid"`
	UserID   string `json:"userId"`
	StreamID string `json:"streamId"`
	Kind     string `json:"kind"`
}

type TrackInfo struct {
	RTPTransceiver *webrtc.RTPTransceiver
	RTPSender      *webrtc.RTPSender
	TrackMetadata  TrackMetadata
}

type trackListener struct {
	log              Logger
	clientID         string
	peerConnection   *webrtc.PeerConnection
	localTracks      []*webrtc.Track
	localTracksMu    sync.RWMutex
	trackInfoByTrack map[*webrtc.Track]TrackInfo
	onTrackEvent     func(TrackEvent)
	mu               sync.RWMutex
	pliInterval      time.Duration
	ssrcMaxBitrates  map[uint32]uint64
}

func newTrackListener(
	loggerFactory LoggerFactory,
	clientID string,
	peerConnection *webrtc.PeerConnection,
	onTrackEvent func(TrackEvent),
) *trackListener {
	p := &trackListener{
		log:              loggerFactory.GetLogger("tracklistener"),
		clientID:         clientID,
		peerConnection:   peerConnection,
		trackInfoByTrack: map[*webrtc.Track]TrackInfo{},
		onTrackEvent:     onTrackEvent,
		ssrcMaxBitrates:  map[uint32]uint64{},
	}

	p.log.Printf("[%s] Setting PeerConnection.OnTrack listener", clientID)
	peerConnection.OnTrack(p.handleTrack)

	return p
}

func (p *trackListener) ClientID() string {
	return p.clientID
}

// GetTracksMetadata gets metadata of the sending tracks with updated Mid
func (p *trackListener) GetTracksMetadata() (metadata []TrackMetadata) {
	p.mu.RLock()
	defer p.mu.RUnlock()

	metadata = make([]TrackMetadata, 0)

	for _, trackInfo := range p.trackInfoByTrack {
		m := trackInfo.TrackMetadata
		m.Mid = trackInfo.RTPTransceiver.Mid()
		metadata = append(metadata, m)
	}
	return
}

func (p *trackListener) WriteRTCP(anyPacket rtcp.Packet) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	switch pkt := anyPacket.(type) {
	case *rtcp.PictureLossIndication:
		prometheusRTCPPacketsSent.Inc()
		return p.peerConnection.WriteRTCP([]rtcp.Packet{pkt})
	case *rtcp.SourceDescription:
	case *rtcp.ReceiverEstimatedMaximumBitrate:
		bitrate, ok := p.ssrcMaxBitrates[pkt.SenderSSRC]
		if !ok || pkt.Bitrate < bitrate {
			p.ssrcMaxBitrates[pkt.SenderSSRC] = bitrate
			prometheusRTCPPacketsSent.Inc()
			return p.peerConnection.WriteRTCP([]rtcp.Packet{pkt})
		}
	case *rtcp.ReceiverReport:
	case *rtcp.SenderReport:
	default:
		p.log.Printf("[%s] Got unhandled RTCP pkt for track: %d (%T)", p.clientID, pkt.DestinationSSRC(), pkt)
	}

	return nil
}

func (p *trackListener) AddTrack(sourceClientID string, track *webrtc.Track) (chan rtcp.Packet, error) {
	p.localTracksMu.Lock()
	defer p.localTracksMu.Unlock()

	p.log.Printf("[%s] peer.AddTrack: %d", p.clientID, track.SSRC())
	rtpSender, err := p.peerConnection.AddTrack(track)

	var transceiver *webrtc.RTPTransceiver
	for _, tr := range p.peerConnection.GetTransceivers() {
		if tr.Sender() == rtpSender {
			transceiver = tr
			break
		}
	}

	rtcpCh := make(chan rtcp.Packet)

	if err != nil {
		close(rtcpCh)
		return rtcpCh, fmt.Errorf("[%s] peer.AddTrack: error adding track: %d: %s", p.clientID, track.SSRC(), err)
	}

	go func() {
		for {
			rtcps, err := rtpSender.ReadRTCP()
			if err != nil {
				p.log.Printf("[%s] RTCP stream for sender track: %d has ended: %s",
					p.clientID,
					track.SSRC(),
					err,
				)
				close(rtcpCh)
				break
			}
			for _, pkt := range rtcps {
				prometheusRTCPPacketsReceived.Inc()
				rtcpCh <- pkt
			}
		}
	}()

	p.trackInfoByTrack[track] = TrackInfo{
		RTPSender:      rtpSender,
		RTPTransceiver: transceiver,
		TrackMetadata: TrackMetadata{
			Mid:      "",
			Kind:     track.Kind().String(),
			UserID:   sourceClientID,
			StreamID: track.Label(),
		},
	}
	return rtcpCh, nil
}

func (p *trackListener) RemoveTrack(track *webrtc.Track) error {
	p.localTracksMu.Lock()
	defer p.localTracksMu.Unlock()
	p.log.Printf("[%s] peer.RemoveTrack: %d", p.clientID, track.SSRC())
	trackInfo, ok := p.trackInfoByTrack[track]
	if !ok {
		return fmt.Errorf("[%s] peer.RemoveTrack: cannot find sender for track: %d", p.clientID, track.SSRC())
	}
	delete(p.trackInfoByTrack, track)
	delete(p.ssrcMaxBitrates, track.SSRC())
	return p.peerConnection.RemoveTrack(trackInfo.RTPSender)
}

func (p *trackListener) handleTrack(remoteTrack *webrtc.Track, receiver *webrtc.RTPReceiver) {
	p.log.Printf("[%s] peer.handleTrack (id: %s, label: %s, type: %s, ssrc: %d)",
		p.clientID, remoteTrack.ID(), remoteTrack.Label(), remoteTrack.Kind(), remoteTrack.SSRC())
	localTrack, err := p.startCopyingTrack(remoteTrack, receiver)
	if err != nil {
		p.log.Printf("Error copying remote track: %s", err)
		return
	}
	p.localTracksMu.Lock()
	p.localTracks = append(p.localTracks, localTrack)
	p.localTracksMu.Unlock()

	p.log.Printf("[%s] peer.handleTrack add track to list of local tracks: %d", p.clientID, localTrack.SSRC())

	p.sendTrackEvent(TrackEvent{p.clientID, localTrack, TrackEventTypeAdd})
}

func (p *trackListener) sendTrackEvent(t TrackEvent) {
	go p.onTrackEvent(t)
}

func (p *trackListener) Tracks() []*webrtc.Track {
	return p.localTracks
}

func (p *trackListener) startCopyingTrack(remoteTrack *webrtc.Track, receiver *webrtc.RTPReceiver) (*webrtc.Track, error) {
	remoteTrackID := remoteTrack.ID()
	if remoteTrackID == "" {
		remoteTrackID = NewUUIDBase62()
	}
	// this is the media stream ID we add the p.clientID in the string to know
	// which user the video came from and the remoteTrack.Label() so we can
	// associate audio/video tracks from the same MediaStream
	remoteTrackLabel := remoteTrack.Label()
	if remoteTrackLabel == "" {
		remoteTrackLabel = NewUUIDBase62()
	}
	localTrackLabel := "sfu_" + p.clientID + "_" + remoteTrackLabel

	localTrackID := "sfu_" + remoteTrackID
	p.log.Printf("[%s] peer.startCopyingTrack: %d", p.clientID, remoteTrack.SSRC())

	ssrc := remoteTrack.SSRC()
	// Create a local track, all our SFU clients will be fed via this track
	localTrack, err := p.peerConnection.NewTrack(remoteTrack.PayloadType(), ssrc, localTrackID, localTrackLabel)
	if err != nil {
		err = fmt.Errorf("[%s] peer.startCopyingTrack: error creating new track: %d, error: %s", p.clientID, remoteTrack.SSRC(), err)
		return nil, err
	}

	// Send a PLI on an interval so that the publisher is pushing a keyframe every rtcpPLIInterval
	// This can be less wasteful by processing incoming RTCP events, then we would emit a NACK/PLI when a viewer requests it

	var ticker *time.Ticker
	if p.pliInterval > 0 {
		ticker = time.NewTicker(p.pliInterval)
		go func() {
			writeRTCP := func() {
				err := p.peerConnection.WriteRTCP(
					[]rtcp.Packet{
						&rtcp.PictureLossIndication{
							MediaSSRC: ssrc,
						},
					},
				)
				if err != nil {
					p.log.Printf("[%s] Error sending rtcp PLI for local track: %d: %s",
						p.clientID,
						localTrack.SSRC(),
						err,
					)
				}
			}
			for range ticker.C {
				writeRTCP()
			}
		}()
	}

	go func() {
		start := time.Now()
		prometheusWebRTCTracksTotal.Inc()
		prometheusWebRTCTracksActive.Inc()
		defer func() {
			prometheusWebRTCTracksActive.Dec()
			prometheusWebRTCTracksDuration.Observe(time.Now().Sub(start).Seconds())
		}()
		defer p.sendTrackEvent(TrackEvent{p.clientID, localTrack, TrackEventTypeRemove})
		if ticker != nil {
			defer ticker.Stop()
		}
		for {
			pkt, err := remoteTrack.ReadRTP()
			if err != nil {
				p.log.Printf("[%s] RTP stream for track: %d has ended: %s", p.clientID, remoteTrack.SSRC(), err)
				return
			}

			prometheusRTPPacketsReceived.Inc()

			// ErrClosedPipe means we don't have any subscribers, this is ok if no peers have connected yet
			err = localTrack.WriteRTP(pkt)
			if err != nil && err != io.ErrClosedPipe {
				p.log.Printf(
					"[%s] Error writing to local track: %d: %s",
					p.clientID,
					localTrack.SSRC(),
					err,
				)
				return
			}
		}
	}()

	return localTrack, nil
}
<|endoftext|>


<commit_before>package service

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

	"github.com/coreos/go-etcd/etcd"
	"github.com/op/go-logging"
)

const (
	servicePrefix  = "service"
	frontEndPrefix = "frontend"
)

type etcdBackend struct {
	client    *etcd.Client
	waitIndex uint64
	Logger    *logging.Logger
	prefix    string
}

func NewEtcdBackend(logger *logging.Logger, uri *url.URL) Backend {
	urls := make([]string, 0)
	if uri.Host != "" {
		urls = append(urls, "http://"+uri.Host)
	}
	return &etcdBackend{
		client: etcd.NewClient(urls),
		prefix: uri.Path,
		Logger: logger,
	}
}

// Watch for changes on a path and return where there is a change.
func (eb *etcdBackend) Watch() error {
	resp, err := eb.client.Watch(eb.prefix, eb.waitIndex, true, nil, nil)
	if err != nil {
		return maskAny(err)
	} else {
		eb.waitIndex = resp.EtcdIndex + 1
		return nil
	}
}

// Load all registered services
func (eb *etcdBackend) Services() ([]ServiceRegistration, error) {
	etcdPath := path.Join(eb.prefix, servicePrefix)
	sort := false
	recursive := true
	resp, err := eb.client.Get(etcdPath, sort, recursive)
	if err != nil {
		return nil, maskAny(err)
	}
	list := []ServiceRegistration{}
	if resp.Node == nil {
		return list, nil
	}
	for _, serviceNode := range resp.Node.Nodes {
		name := path.Base(serviceNode.Key)
		registrations := make(map[int]*ServiceRegistration)
		for _, backendNode := range serviceNode.Nodes {
			uniqueID := path.Base(backendNode.Key)
			parts := strings.Split(uniqueID, ":")
			if len(parts) < 3 {
				eb.Logger.Warning("UniqueID malformed: '%s'", uniqueID)
				continue
			}
			port, err := strconv.Atoi(parts[1])
			if err != nil {
				continue
			}
			sr, ok := registrations[port]
			if !ok {
				sr = &ServiceRegistration{ServiceName: name, Port: port}
				registrations[port] = sr
			}
			sr.Backends = append(sr.Backends, backendNode.Value)
		}
		for _, v := range registrations {
			list = append(list, *v)
		}
	}

	return list, nil
}

type frontendRecord struct {
	Selectors     []frontendSelectorRecord `json:"selectors"`
	Service       string                   `json:"service,omitempty"`
	HttpCheckPath string                   `json:"http-check-path,omitempty"`
}

type frontendSelectorRecord struct {
	Domain     string `json:"domain,omitempty"`
	SslCert    string `json:"ssl-cert,omitempty"`
	PathPrefix string `json:"path-prefix,omitempty"`
	Port       int    `json:"port,omitempty"`
	Private    bool   `json:"private,omitempty"`
}

// Load all registered front-ends
func (eb *etcdBackend) FrontEnds() ([]FrontEndRegistration, error) {
	etcdPath := path.Join(eb.prefix, frontEndPrefix)
	sort := false
	recursive := false
	resp, err := eb.client.Get(etcdPath, sort, recursive)
	if err != nil {
		return nil, maskAny(err)
	}
	list := []FrontEndRegistration{}
	if resp.Node == nil {
		return list, nil
	}
	for _, frontEndNode := range resp.Node.Nodes {
		rawJson := frontEndNode.Value
		record := &frontendRecord{}
		if err := json.Unmarshal([]byte(rawJson), record); err != nil {
			eb.Logger.Error("Cannot unmarshal registration of %s", frontEndNode.Key)
			continue
		}

		name := path.Base(frontEndNode.Key)
		registrations := make(map[int]*FrontEndRegistration)
		for _, sel := range record.Selectors {
			port := sel.Port
			reg, ok := registrations[port]
			if !ok {
				reg := &FrontEndRegistration{
					Name:          name,
					Service:       record.Service,
					Port:          port,
					HttpCheckPath: record.HttpCheckPath,
				}
				registrations[port] = reg
			}
			reg.Selectors = append(reg.Selectors, FrontEndSelector{
				Domain:     sel.Domain,
				SslCert:    sel.SslCert,
				PathPrefix: sel.PathPrefix,
				Port:       sel.Port,
				Private:    sel.Private,
			})
		}
		for _, reg := range registrations {
			list = append(list, *reg)
		}
	}

	return list, nil
}
<commit_msg>Fix crash<commit_after>package service

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

	"github.com/coreos/go-etcd/etcd"
	"github.com/op/go-logging"
)

const (
	servicePrefix  = "service"
	frontEndPrefix = "frontend"
)

type etcdBackend struct {
	client    *etcd.Client
	waitIndex uint64
	Logger    *logging.Logger
	prefix    string
}

func NewEtcdBackend(logger *logging.Logger, uri *url.URL) Backend {
	urls := make([]string, 0)
	if uri.Host != "" {
		urls = append(urls, "http://"+uri.Host)
	}
	return &etcdBackend{
		client: etcd.NewClient(urls),
		prefix: uri.Path,
		Logger: logger,
	}
}

// Watch for changes on a path and return where there is a change.
func (eb *etcdBackend) Watch() error {
	resp, err := eb.client.Watch(eb.prefix, eb.waitIndex, true, nil, nil)
	if err != nil {
		return maskAny(err)
	} else {
		eb.waitIndex = resp.EtcdIndex + 1
		return nil
	}
}

// Load all registered services
func (eb *etcdBackend) Services() ([]ServiceRegistration, error) {
	etcdPath := path.Join(eb.prefix, servicePrefix)
	sort := false
	recursive := true
	resp, err := eb.client.Get(etcdPath, sort, recursive)
	if err != nil {
		return nil, maskAny(err)
	}
	list := []ServiceRegistration{}
	if resp.Node == nil {
		return list, nil
	}
	for _, serviceNode := range resp.Node.Nodes {
		name := path.Base(serviceNode.Key)
		registrations := make(map[int]*ServiceRegistration)
		for _, backendNode := range serviceNode.Nodes {
			uniqueID := path.Base(backendNode.Key)
			parts := strings.Split(uniqueID, ":")
			if len(parts) < 3 {
				eb.Logger.Warning("UniqueID malformed: '%s'", uniqueID)
				continue
			}
			port, err := strconv.Atoi(parts[1])
			if err != nil {
				continue
			}
			sr, ok := registrations[port]
			if !ok {
				sr = &ServiceRegistration{ServiceName: name, Port: port}
				registrations[port] = sr
			}
			sr.Backends = append(sr.Backends, backendNode.Value)
		}
		for _, v := range registrations {
			list = append(list, *v)
		}
	}

	return list, nil
}

type frontendRecord struct {
	Selectors     []frontendSelectorRecord `json:"selectors"`
	Service       string                   `json:"service,omitempty"`
	HttpCheckPath string                   `json:"http-check-path,omitempty"`
}

type frontendSelectorRecord struct {
	Domain     string `json:"domain,omitempty"`
	SslCert    string `json:"ssl-cert,omitempty"`
	PathPrefix string `json:"path-prefix,omitempty"`
	Port       int    `json:"port,omitempty"`
	Private    bool   `json:"private,omitempty"`
}

// Load all registered front-ends
func (eb *etcdBackend) FrontEnds() ([]FrontEndRegistration, error) {
	etcdPath := path.Join(eb.prefix, frontEndPrefix)
	sort := false
	recursive := false
	resp, err := eb.client.Get(etcdPath, sort, recursive)
	if err != nil {
		return nil, maskAny(err)
	}
	list := []FrontEndRegistration{}
	if resp.Node == nil {
		return list, nil
	}
	for _, frontEndNode := range resp.Node.Nodes {
		rawJson := frontEndNode.Value
		record := &frontendRecord{}
		if err := json.Unmarshal([]byte(rawJson), record); err != nil {
			eb.Logger.Error("Cannot unmarshal registration of %s", frontEndNode.Key)
			continue
		}

		name := path.Base(frontEndNode.Key)
		registrations := make(map[int]*FrontEndRegistration)
		for _, sel := range record.Selectors {
			port := sel.Port
			reg, ok := registrations[port]
			if !ok {
				reg = &FrontEndRegistration{
					Name:          name,
					Service:       record.Service,
					Port:          port,
					HttpCheckPath: record.HttpCheckPath,
				}
				registrations[port] = reg
			}
			reg.Selectors = append(reg.Selectors, FrontEndSelector{
				Domain:     sel.Domain,
				SslCert:    sel.SslCert,
				PathPrefix: sel.PathPrefix,
				Port:       sel.Port,
				Private:    sel.Private,
			})
		}
		for _, reg := range registrations {
			list = append(list, *reg)
		}
	}

	return list, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/sdwolfe32/ANIRip/anirip"
)

// Trims the first couple seconds off of the video to remove any logos
func trimMKV(adLength int, engineDir, tempDir string) error {
	// Removes a stale temp files to avoid conflcts in func
	os.Remove(tempDir + "\\" + "untrimmed.episode.mkv")
	os.Remove(tempDir + "\\" + "split.episode-001.mkv")
	os.Remove(tempDir + "\\" + "prefix.episode.mkv")
	os.Remove(tempDir + "\\" + "split.episode-002.mkv")
	os.Remove(tempDir + "\\" + "list.episode.txt")

	// Recursively retries rename to temp filename before execution
	if err := anirip.Rename(tempDir+"\\episode.mkv", tempDir+"\\untrimmed.episode.mkv", 10); err != nil {
		return err
	}

	// Finds the clis we need for trimming
	ffmpeg, err := filepath.Abs(engineDir + "\\ffmpeg.exe")
	if err != nil {
		return anirip.Error{Message: "Unable to find ffmpeg.exe in \\" + engineDir + "\\ directory", Err: err}
	}
	mkvmerge, err := filepath.Abs(engineDir + "\\mkvmerge.exe")
	if err != nil {
		return anirip.Error{Message: "Unable to find ffmpeg.exe in \\" + engineDir + "\\ directory", Err: err}
	}

	// Creates the command too split the meat of the video from the first ad chunk
	cmd := exec.Command(mkvmerge,
		"--split", "timecodes:"+anirip.MStoTimecode(adLength),
		"-o", "split.episode.mkv",
		"untrimmed.episode.mkv",
	)
	cmd.Dir = tempDir // Sets working directory to temp so our halves end up there

	// Executes the command
	_, err = cmd.Output()
	if err != nil {
		return anirip.Error{Message: "There was an error while splitting the episode", Err: err}
	}

	// Executes the fine intro trim and waits for the command to finish
	cmd = exec.Command(ffmpeg,
		"-i", "split.episode-001.mkv",
		"-ss", anirip.MStoTimecode(adLength), // Exact timestamp of the ad endings
		"-c:v", "libx264",
		"-c:a", "aac",
		"-preset", "slow",
		"-crf", "5", "-y", // Use AAC as audio codec to match video.mkv
		"prefix.episode.mkv")
	cmd.Dir = tempDir // Sets working directory to temp

	// Executes the command
	_, err = cmd.Output()
	if err != nil {
		return anirip.Error{Message: "There was an error while creating the prefix clip", Err: err}
	}

	// Creates a text file containing the file names of the 2 files created above
	fileListBytes := []byte("file 'prefix.episode.mkv'\r\nfile 'split.episode-002.mkv'")
	if err = ioutil.WriteFile(tempDir+"\\"+"list.episode.txt", fileListBytes, 0644); err != nil {
		return anirip.Error{Message: "There was an error while creating list.episode.txt", Err: err}
	}

	// Executes the merge of our two temporary files
	cmd = exec.Command(ffmpeg,
		"-f", "concat",
		"-i", "list.episode.txt",
		"-c", "copy", "-y",
		"episode.mkv")
	cmd.Dir = tempDir // Sets working directory to temp

	// Executes the command
	_, err = cmd.Output()
	if err != nil {
		return anirip.Error{Message: "There was an error while merging video and prefix", Err: err}
	}

	// Removes the temporary files we created as they are no longer needed
	os.Remove(tempDir + "\\" + "untrimmed.episode.mkv")
	os.Remove(tempDir + "\\" + "split.episode-001.mkv")
	os.Remove(tempDir + "\\" + "prefix.episode.mkv")
	os.Remove(tempDir + "\\" + "split.episode-002.mkv")
	os.Remove(tempDir + "\\" + "list.episode.txt")
	return nil
}

// Merges a VIDEO.mkv and a VIDEO.ass
func mergeSubtitles(audioLang, subtitleLang, engineDir, tempDir string) error {
	// Removes a stale temp files to avoid conflcts in func
	os.Remove(tempDir + "\\unmerged.episode.mkv")

	// Recursively retries rename to temp filename before execution
	if err := anirip.Rename(tempDir+"\\episode.mkv", tempDir+"\\unmerged.episode.mkv", 10); err != nil {
		return err
	}

	path, err := filepath.Abs(engineDir + "\\ffmpeg.exe")
	if err != nil {
		return anirip.Error{Message: "Unable to find ffmpeg.exe in \\" + engineDir + "\\ directory", Err: err}
	}

	// Creates the command which we will use to merge our subtitles and video
	cmd := new(exec.Cmd)
	if subtitleLang == "" {
		cmd = exec.Command(path,
			"-i", "unmerged.episode.mkv",
			"-c:v", "copy",
			"-c:a", "copy",
			"-metadata:s:a:0", "language="+audioLang, // sets audio language to passed audioLang
			"-y", "episode.mkv")
	} else {
		cmd = exec.Command(path,
			"-i", "unmerged.episode.mkv",
			"-f", "ass",
			"-i", "subtitles.episode.ass",
			"-c:v", "copy",
			"-c:a", "copy",
			"-metadata:s:a:0", "language="+audioLang, // sets audio language to passed audioLang
			"-metadata:s:s:0", "language="+subtitleLang, // sets subtitle language to subtitleLang
			"-disposition:s:0", "default",
			"-y", "episode.mkv")
	}
	cmd.Dir = tempDir // Sets working directory to temp

	// Executes the command
	_, err = cmd.Output()
	if err != nil {
		return anirip.Error{Message: "There was an error while merging subtitles", Err: err}
	}

	// Removes old temp files
	os.Remove(tempDir + "\\subtitles.episode.ass")
	os.Remove(tempDir + "\\unmerged.episode.mkv")
	return nil
}

// Cleans up the mkv, optimizing it for playback
func cleanMKV(engineDir, tempDir string) error {
	// Removes a stale temp file to avoid conflcts in func
	os.Remove(tempDir + "\\dirty.episode.mkv")

	// Recursively retries rename to temp filename before execution
	if err := anirip.Rename(tempDir+"\\episode.mkv", tempDir+"\\"+"dirty.episode.mkv", 10); err != nil {
		return err
	}

	// Finds the path of mkclean.exe so we can perform system calls on it
	path, err := filepath.Abs(engineDir + "\\mkclean.exe")
	if err != nil {
		return anirip.Error{Message: "Unable to find mkclean.exe in \\" + engineDir + "\\ directory", Err: err}
	}

	// Creates the command which we will use to clean our mkv to "video.clean.mkv"
	cmd := exec.Command(path,
		"--optimize",
		"dirty.episode.mkv",
		"episode.mkv")
	cmd.Dir = tempDir // Sets working directory to temp

	// Executes the command
	_, err = cmd.Output()
	if err != nil {
		return anirip.Error{Message: "There was an error while optimizing our mkv", Err: err}
	}

	// Deletes the old, un-needed dirty mkv file
	os.Remove(tempDir + "\\dirty.episode.mkv")
	return nil
}

// Gets user input from the user and unmarshalls it into the input
func getStandardUserInput(prefixText string, input *string) error {
	fmt.Printf(prefixText)
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		*input = scanner.Text()
		break
	}
	if err := scanner.Err(); err != nil {
		return anirip.Error{Message: "There was an error getting standard user input", Err: err}
	}
	return nil
}

// Blocks execution and waits for the user to press enter
func pause() {
	fmt.Print("Press 'Enter' to continue...")
	bufio.NewReader(os.Stdin).ReadBytes('\n')
}
<commit_msg>Flawless audio/video transition by just copying audio stream and rencoding only first few seconds of video stream<commit_after>package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/sdwolfe32/ANIRip/anirip"
)

// Trims the first couple seconds off of the video to remove any logos
func trimMKV(adLength int, engineDir, tempDir string) error {
	// Removes a stale temp files to avoid conflcts in func
	os.Remove(tempDir + "\\" + "untrimmed.episode.mkv")
	os.Remove(tempDir + "\\" + "split.episode-001.mkv")
	os.Remove(tempDir + "\\" + "prefix.episode.mkv")
	os.Remove(tempDir + "\\" + "split.episode-002.mkv")
	os.Remove(tempDir + "\\" + "list.episode.txt")

	// Recursively retries rename to temp filename before execution
	if err := anirip.Rename(tempDir+"\\episode.mkv", tempDir+"\\untrimmed.episode.mkv", 10); err != nil {
		return err
	}

	// Finds the clis we need for trimming
	ffmpeg, err := filepath.Abs(engineDir + "\\ffmpeg.exe")
	if err != nil {
		return anirip.Error{Message: "Unable to find ffmpeg.exe in \\" + engineDir + "\\ directory", Err: err}
	}
	mkvmerge, err := filepath.Abs(engineDir + "\\mkvmerge.exe")
	if err != nil {
		return anirip.Error{Message: "Unable to find ffmpeg.exe in \\" + engineDir + "\\ directory", Err: err}
	}

	// Creates the command too split the meat of the video from the first ad chunk
	cmd := exec.Command(mkvmerge,
		"--split", "timecodes:"+anirip.MStoTimecode(adLength),
		"-o", "split.episode.mkv",
		"untrimmed.episode.mkv",
	)
	cmd.Dir = tempDir // Sets working directory to temp so our halves end up there

	// Executes the command
	_, err = cmd.Output()
	if err != nil {
		return anirip.Error{Message: "There was an error while splitting the episode", Err: err}
	}

	// Executes the fine intro trim and waits for the command to finish
	cmd = exec.Command(ffmpeg,
		"-i", "split.episode-001.mkv",
		"-ss", anirip.MStoTimecode(adLength), // Exact timestamp of the ad endings
		"-c:v", "h264",
		"-crf", "15",
		"-preset", "slow",
		"-c:a", "copy", "-y", // Use AAC as audio codec to match video.mkv
		"prefix.episode.mkv")
	cmd.Dir = tempDir // Sets working directory to temp

	// Executes the command
	_, err = cmd.Output()
	if err != nil {
		return anirip.Error{Message: "There was an error while creating the prefix clip", Err: err}
	}

	// Creates a text file containing the file names of the 2 files created above
	fileListBytes := []byte("file 'prefix.episode.mkv'\r\nfile 'split.episode-002.mkv'")
	if err = ioutil.WriteFile(tempDir+"\\"+"list.episode.txt", fileListBytes, 0644); err != nil {
		return anirip.Error{Message: "There was an error while creating list.episode.txt", Err: err}
	}

	// Executes the merge of our two temporary files
	cmd = exec.Command(ffmpeg,
		"-f", "concat",
		"-i", "list.episode.txt",
		"-c", "copy", "-y",
		"episode.mkv")
	cmd.Dir = tempDir // Sets working directory to temp

	// Executes the command
	_, err = cmd.Output()
	if err != nil {
		return anirip.Error{Message: "There was an error while merging video and prefix", Err: err}
	}

	// Removes the temporary files we created as they are no longer needed
	os.Remove(tempDir + "\\" + "untrimmed.episode.mkv")
	os.Remove(tempDir + "\\" + "split.episode-001.mkv")
	os.Remove(tempDir + "\\" + "prefix.episode.mkv")
	os.Remove(tempDir + "\\" + "split.episode-002.mkv")
	os.Remove(tempDir + "\\" + "list.episode.txt")
	return nil
}

// Merges a VIDEO.mkv and a VIDEO.ass
func mergeSubtitles(audioLang, subtitleLang, engineDir, tempDir string) error {
	// Removes a stale temp files to avoid conflcts in func
	os.Remove(tempDir + "\\unmerged.episode.mkv")

	// Recursively retries rename to temp filename before execution
	if err := anirip.Rename(tempDir+"\\episode.mkv", tempDir+"\\unmerged.episode.mkv", 10); err != nil {
		return err
	}

	path, err := filepath.Abs(engineDir + "\\ffmpeg.exe")
	if err != nil {
		return anirip.Error{Message: "Unable to find ffmpeg.exe in \\" + engineDir + "\\ directory", Err: err}
	}

	// Creates the command which we will use to merge our subtitles and video
	cmd := new(exec.Cmd)
	if subtitleLang == "" {
		cmd = exec.Command(path,
			"-i", "unmerged.episode.mkv",
			"-c:v", "copy",
			"-c:a", "copy",
			"-metadata:s:a:0", "language="+audioLang, // sets audio language to passed audioLang
			"-y", "episode.mkv")
	} else {
		cmd = exec.Command(path,
			"-i", "unmerged.episode.mkv",
			"-f", "ass",
			"-i", "subtitles.episode.ass",
			"-c:v", "copy",
			"-c:a", "copy",
			"-metadata:s:a:0", "language="+audioLang, // sets audio language to passed audioLang
			"-metadata:s:s:0", "language="+subtitleLang, // sets subtitle language to subtitleLang
			"-disposition:s:0", "default",
			"-y", "episode.mkv")
	}
	cmd.Dir = tempDir // Sets working directory to temp

	// Executes the command
	_, err = cmd.Output()
	if err != nil {
		return anirip.Error{Message: "There was an error while merging subtitles", Err: err}
	}

	// Removes old temp files
	os.Remove(tempDir + "\\subtitles.episode.ass")
	os.Remove(tempDir + "\\unmerged.episode.mkv")
	return nil
}

// Cleans up the mkv, optimizing it for playback
func cleanMKV(engineDir, tempDir string) error {
	// Removes a stale temp file to avoid conflcts in func
	os.Remove(tempDir + "\\dirty.episode.mkv")

	// Recursively retries rename to temp filename before execution
	if err := anirip.Rename(tempDir+"\\episode.mkv", tempDir+"\\"+"dirty.episode.mkv", 10); err != nil {
		return err
	}

	// Finds the path of mkclean.exe so we can perform system calls on it
	path, err := filepath.Abs(engineDir + "\\mkclean.exe")
	if err != nil {
		return anirip.Error{Message: "Unable to find mkclean.exe in \\" + engineDir + "\\ directory", Err: err}
	}

	// Creates the command which we will use to clean our mkv to "video.clean.mkv"
	cmd := exec.Command(path,
		"--optimize",
		"dirty.episode.mkv",
		"episode.mkv")
	cmd.Dir = tempDir // Sets working directory to temp

	// Executes the command
	_, err = cmd.Output()
	if err != nil {
		return anirip.Error{Message: "There was an error while optimizing our mkv", Err: err}
	}

	// Deletes the old, un-needed dirty mkv file
	os.Remove(tempDir + "\\dirty.episode.mkv")
	return nil
}

// Gets user input from the user and unmarshalls it into the input
func getStandardUserInput(prefixText string, input *string) error {
	fmt.Printf(prefixText)
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		*input = scanner.Text()
		break
	}
	if err := scanner.Err(); err != nil {
		return anirip.Error{Message: "There was an error getting standard user input", Err: err}
	}
	return nil
}

// Blocks execution and waits for the user to press enter
func pause() {
	fmt.Print("Press 'Enter' to continue...")
	bufio.NewReader(os.Stdin).ReadBytes('\n')
}
<|endoftext|>


<commit_before>package kasia

import (
    "io"
    "io/ioutil"
    "os"
    "bytes"
)

// Funkcje i metody kompatybilne z Go template.

func (tpl *Template) Execute(data interface{}, wr io.Writer) os.Error {
    return tpl.Run(wr, data)
}

func (tpl *Template) ParseFile(filename string) (err os.Error) {
    data, err := ioutil.ReadFile(filename)
    return tpl.Parse(string(data))
}

func Parse(str string) (tpl *Template, err os.Error) {
    tpl = New()
    err = tpl.Parse(str)
    if err != nil {
        tpl = nil
    }
    return
}

func ParseFile(filename string) (tpl *Template, err os.Error) {
    tpl = New()
    err = tpl.ParseFile(filename)
    if err != nil {
        tpl = nil
    }
    return
}

// Funkcje i metody kompatybilne z mustache.go

func (tpl *Template) Render(ctx ...interface{}) string {
    var buf bytes.Buffer
    err := tpl.Run(&buf, ctx...)
    if err != nil {
        panic(err)
    }
    return buf.String()
}
<commit_msg>Added RenderTxt method<commit_after>package kasia

import (
    "io"
    "io/ioutil"
    "os"
    "bytes"
)

// Renderowanie do tekstu - metoda przydatna przy testowaniu

func (tpl *Template) RenderTxt(ctx ...interface{}) (string, os.Error) {
    var buf bytes.Buffer
    err := tpl.Run(&buf, ctx...)
    if err != nil {
        return "", err
    }
    return buf.String(), nil
}

// Funkcje i metody kompatybilne z Go template

func (tpl *Template) Execute(data interface{}, wr io.Writer) os.Error {
    return tpl.Run(wr, data)
}

func (tpl *Template) ParseFile(filename string) (err os.Error) {
    data, err := ioutil.ReadFile(filename)
    return tpl.Parse(string(data))
}

func Parse(str string) (tpl *Template, err os.Error) {
    tpl = New()
    err = tpl.Parse(str)
    if err != nil {
        tpl = nil
    }
    return
}

func ParseFile(filename string) (tpl *Template, err os.Error) {
    tpl = New()
    err = tpl.ParseFile(filename)
    if err != nil {
        tpl = nil
    }
    return
}

// Funkcje i metody kompatybilne z mustache.go

func (tpl *Template) Render(ctx ...interface{}) (out string) {
    var err os.Error
    out, err = tpl.RenderTxt(ctx...)
    if err != nil {
        panic(err)
    }
    return
}

<|endoftext|>


<commit_before>package writeaheadlog

const (
	checksumSize   = 16
	pageSize       = 4096
	pageMetaSize   = checksumSize + 4*8 // checksum + 4 uint64s

	// MaxPayloadSize is the number of bytes that can fit into a single
	// page. For best performance, the number of pages written should be
	// minimized, so clients should try to keep the length of an Update's
	// Instructions field slightly below a multiple of MaxPayloadSize.
	MaxPayloadSize = pageSize - pageMetaSize
)

const (
	pageStatusInvalid = iota
	pageStatusOther
	pageStatusWritten
	pageStatusComitted
	pageStatusApplied
)

const (
	recoveryStateInvalid = iota
	recoveryStateClean
	recoveryStateUnclean
	recoveryStateWipe
)

const (
	metadataHeader  = "WAL"
	metadataVersion = "1.0"
)

// A checksum is a 128-bit blake2b hash.
type checksum [checksumSize]byte
<commit_msg>go fmt<commit_after>package writeaheadlog

const (
	checksumSize = 16
	pageSize     = 4096
	pageMetaSize = checksumSize + 4*8 // checksum + 4 uint64s

	// MaxPayloadSize is the number of bytes that can fit into a single
	// page. For best performance, the number of pages written should be
	// minimized, so clients should try to keep the length of an Update's
	// Instructions field slightly below a multiple of MaxPayloadSize.
	MaxPayloadSize = pageSize - pageMetaSize
)

const (
	pageStatusInvalid = iota
	pageStatusOther
	pageStatusWritten
	pageStatusComitted
	pageStatusApplied
)

const (
	recoveryStateInvalid = iota
	recoveryStateClean
	recoveryStateUnclean
	recoveryStateWipe
)

const (
	metadataHeader  = "WAL"
	metadataVersion = "1.0"
)

// A checksum is a 128-bit blake2b hash.
type checksum [checksumSize]byte
<|endoftext|>


<commit_before>package boom

import (
	"bytes"
	"encoding/binary"
	"errors"
	"hash"
	"hash/fnv"
	"math"
	"math/rand"
)

// maxNumKicks is the maximum number of relocations to attempt when inserting
// an element before considering the filter full.
const maxNumKicks = 500

// bucket consists of a set of []byte entries.
type bucket [][]byte

// contains indicates if the given fingerprint is contained in one of the
// bucket's entries.
func (b bucket) contains(f []byte) bool {
	return b.indexOf(f) != -1
}

// indexOf returns the entry index of the given fingerprint or -1 if it's not
// in the bucket.
func (b bucket) indexOf(f []byte) int {
	for i, fingerprint := range b {
		if bytes.Equal(f, fingerprint) {
			return i
		}
	}
	return -1
}

// getEmptyEntry returns the index of the next available entry in the bucket or
// an error if it's full.
func (b bucket) getEmptyEntry() (int, error) {
	for i, fingerprint := range b {
		if fingerprint == nil {
			return i, nil
		}
	}
	return -1, errors.New("full")
}

// CuckooFilter implements a Cuckoo Bloom filter as described by Andersen,
// Kaminsky, and Mitzenmacher in Cuckoo Filter: Practically Better Than Bloom:
//
// http://www.pdl.cmu.edu/PDL-FTP/FS/cuckoo-conext2014.pdf
//
// A Cuckoo Filter is a Bloom filter variation which provides support for
// removing elements without significantly degrading space and performance. It
// works by using a cuckoo hashing scheme for inserting items. Instead of
// storing the elements themselves, it stores their fingerprints which also
// allows for item removal without false negatives (if you don't attempt to
// remove an item not contained in the filter).
//
// For applications that store many items and target moderately low
// false-positive rates, cuckoo filters have lower space overhead than
// space-optimized Bloom filters.
type CuckooFilter struct {
	buckets []bucket
	hash    hash.Hash32 // hash function (used for fingerprint and hash)
	m       uint        // number of buckets
	b       uint        // number of entries per bucket
	f       uint        // length of fingerprints (in bytes)
	count   uint        // number of items in the filter
	n       uint        // filter capacity
}

// NewCuckooFilter creates a new Cuckoo Bloom filter optimized to store n items
// with a specified target false-positive rate.
func NewCuckooFilter(n uint, fpRate float64) *CuckooFilter {
	var (
		epsilon = 1 - fpRate
		b       = uint(4)
		f       = calculateF(b, epsilon)
		m       = power2(n / uint(f) * 8)
		buckets = make([]bucket, m)
	)

	for i := uint(0); i < m; i++ {
		buckets[i] = make(bucket, b)
	}

	return &CuckooFilter{
		buckets: buckets,
		hash:    fnv.New32(),
		m:       m,
		b:       b,
		f:       uint(f),
		n:       n,
	}
}

// Buckets returns the number of buckets.
func (c *CuckooFilter) Buckets() uint {
	return c.m
}

// Capacity returns the number of items the filter can store.
func (c *CuckooFilter) Capacity() uint {
	return c.n
}

// Count returns the number of items in the filter.
func (c *CuckooFilter) Count() uint {
	return c.count
}

// Test will test for membership of the data and returns true if it is a
// member, false if not. This is a probabilistic test, meaning there is a
// non-zero probability of false positives.
func (c *CuckooFilter) Test(data []byte) bool {
	i1, i2, f := c.components(data)

	// If either bucket contains f, it's a member.
	return c.buckets[i1%c.m].contains(f) || c.buckets[i2%c.m].contains(f)
}

// Add will add the data to the Cuckoo Filter. It returns an error if the
// filter is full. If the filter is full, an item is removed to make room for
// the new item. This introduces a possibility for false negatives. To avoid
// this, use Count and Capacity to check if the filter is full before adding an
// item.
func (c *CuckooFilter) Add(data []byte) error {
	return c.add(c.components(data))
}

// TestAndAdd is equivalent to calling Test followed by Add. It returns true if
// the data is a member, false if not. An error is returned if the filter is
// full. If the filter is full, an item is removed to make room for the new
// item. This introduces a possibility for false negatives. To avoid this, use
// Count and Capacity to check if the filter is full before adding an item.
func (c *CuckooFilter) TestAndAdd(data []byte) (bool, error) {
	i1, i2, f := c.components(data)

	// If either bucket contains f, it's a member.
	if c.buckets[i1%c.m].contains(f) || c.buckets[i2%c.m].contains(f) {
		return true, nil
	}

	return false, c.add(i1, i2, f)
}

// TestAndRemove will test for membership of the data and remove it from the
// filter if it exists. Returns true if the data was a member, false if not.
func (c *CuckooFilter) TestAndRemove(data []byte) bool {
	i1, i2, f := c.components(data)

	// Try to remove from bucket[i1].
	b1 := c.buckets[i1%c.m]
	if idx := b1.indexOf(f); idx != -1 {
		b1[idx] = nil
		c.count--
		return true
	}

	// Try to remove from bucket[i2].
	b2 := c.buckets[i2%c.m]
	if idx := b2.indexOf(f); idx != -1 {
		b2[idx] = nil
		c.count--
		return true
	}

	return false
}

// Reset restores the Bloom filter to its original state. It returns the filter
// to allow for chaining.
func (c *CuckooFilter) Reset() *CuckooFilter {
	buckets := make([]bucket, c.m)
	for i := uint(0); i < c.m; i++ {
		buckets[i] = make(bucket, c.b)
	}
	c.buckets = buckets
	c.count = 0
	return c
}

// add will insert the fingerprint into the filter returning an error if the
// filter is full.
func (c *CuckooFilter) add(i1, i2 uint, f []byte) error {
	// Try to insert into bucket[i1].
	b1 := c.buckets[i1%c.m]
	if idx, err := b1.getEmptyEntry(); err == nil {
		b1[idx] = f
		c.count++
		return nil
	}

	// Try to insert into bucket[i2].
	b2 := c.buckets[i2%c.m]
	if idx, err := b2.getEmptyEntry(); err == nil {
		b2[idx] = f
		c.count++
		return nil
	}

	// Must relocate existing items.
	i := i1
	for n := 0; n < maxNumKicks; n++ {
		idx := i % c.m
		tmp := c.buckets[idx][rand.Intn(int(c.b))]
		c.buckets[idx][0] = f
		f = tmp
		i = i ^ uint(binary.BigEndian.Uint32(c.computeHash(f)))
		b := c.buckets[i%c.m]
		if idx, err := b.getEmptyEntry(); err == nil {
			b[idx] = f
			c.count++
			return nil
		}
	}

	return errors.New("full")
}

// components returns the two hash values used to index into the buckets and
// the fingerprint for the given element.
func (c *CuckooFilter) components(data []byte) (uint, uint, []byte) {
	var (
		hash = c.computeHash(data)
		f    = hash[0:c.f]
		i1   = uint(binary.BigEndian.Uint32(hash))
		i2   = i1 ^ uint(binary.BigEndian.Uint32(c.computeHash(f)))
	)

	return i1, i2, f
}

// computeHash returns a 32-bit hash value for the given data.
func (c *CuckooFilter) computeHash(data []byte) []byte {
	c.hash.Write(data)
	hash := c.hash.Sum(nil)
	c.hash.Reset()
	return hash
}

// calculateF returns the optimal fingerprint length in bytes for the given
// bucket size and false-positive rate epsilon.
func calculateF(b uint, epsilon float64) uint {
	f := uint(math.Ceil(math.Log(2 * float64(b) / epsilon)))
	f = f / 8
	if f <= 0 {
		f = 1
	}
	return f
}

// power2 calculates the next power of two for the given value.
func power2(x uint) uint {
	x--
	x |= x >> 1
	x |= x >> 2
	x |= x >> 4
	x |= x >> 8
	x |= x >> 16
	x |= x >> 32
	x++
	return x
}
<commit_msg>Fix epsilon error<commit_after>package boom

import (
	"bytes"
	"encoding/binary"
	"errors"
	"hash"
	"hash/fnv"
	"math"
	"math/rand"
)

// maxNumKicks is the maximum number of relocations to attempt when inserting
// an element before considering the filter full.
const maxNumKicks = 500

// bucket consists of a set of []byte entries.
type bucket [][]byte

// contains indicates if the given fingerprint is contained in one of the
// bucket's entries.
func (b bucket) contains(f []byte) bool {
	return b.indexOf(f) != -1
}

// indexOf returns the entry index of the given fingerprint or -1 if it's not
// in the bucket.
func (b bucket) indexOf(f []byte) int {
	for i, fingerprint := range b {
		if bytes.Equal(f, fingerprint) {
			return i
		}
	}
	return -1
}

// getEmptyEntry returns the index of the next available entry in the bucket or
// an error if it's full.
func (b bucket) getEmptyEntry() (int, error) {
	for i, fingerprint := range b {
		if fingerprint == nil {
			return i, nil
		}
	}
	return -1, errors.New("full")
}

// CuckooFilter implements a Cuckoo Bloom filter as described by Andersen,
// Kaminsky, and Mitzenmacher in Cuckoo Filter: Practically Better Than Bloom:
//
// http://www.pdl.cmu.edu/PDL-FTP/FS/cuckoo-conext2014.pdf
//
// A Cuckoo Filter is a Bloom filter variation which provides support for
// removing elements without significantly degrading space and performance. It
// works by using a cuckoo hashing scheme for inserting items. Instead of
// storing the elements themselves, it stores their fingerprints which also
// allows for item removal without false negatives (if you don't attempt to
// remove an item not contained in the filter).
//
// For applications that store many items and target moderately low
// false-positive rates, cuckoo filters have lower space overhead than
// space-optimized Bloom filters.
type CuckooFilter struct {
	buckets []bucket
	hash    hash.Hash32 // hash function (used for fingerprint and hash)
	m       uint        // number of buckets
	b       uint        // number of entries per bucket
	f       uint        // length of fingerprints (in bytes)
	count   uint        // number of items in the filter
	n       uint        // filter capacity
}

// NewCuckooFilter creates a new Cuckoo Bloom filter optimized to store n items
// with a specified target false-positive rate.
func NewCuckooFilter(n uint, fpRate float64) *CuckooFilter {
	var (
		b       = uint(4)
		f       = calculateF(b, fpRate)
		m       = power2(n / uint(f) * 8)
		buckets = make([]bucket, m)
	)

	for i := uint(0); i < m; i++ {
		buckets[i] = make(bucket, b)
	}

	return &CuckooFilter{
		buckets: buckets,
		hash:    fnv.New32(),
		m:       m,
		b:       b,
		f:       uint(f),
		n:       n,
	}
}

// Buckets returns the number of buckets.
func (c *CuckooFilter) Buckets() uint {
	return c.m
}

// Capacity returns the number of items the filter can store.
func (c *CuckooFilter) Capacity() uint {
	return c.n
}

// Count returns the number of items in the filter.
func (c *CuckooFilter) Count() uint {
	return c.count
}

// Test will test for membership of the data and returns true if it is a
// member, false if not. This is a probabilistic test, meaning there is a
// non-zero probability of false positives.
func (c *CuckooFilter) Test(data []byte) bool {
	i1, i2, f := c.components(data)

	// If either bucket contains f, it's a member.
	return c.buckets[i1%c.m].contains(f) || c.buckets[i2%c.m].contains(f)
}

// Add will add the data to the Cuckoo Filter. It returns an error if the
// filter is full. If the filter is full, an item is removed to make room for
// the new item. This introduces a possibility for false negatives. To avoid
// this, use Count and Capacity to check if the filter is full before adding an
// item.
func (c *CuckooFilter) Add(data []byte) error {
	return c.add(c.components(data))
}

// TestAndAdd is equivalent to calling Test followed by Add. It returns true if
// the data is a member, false if not. An error is returned if the filter is
// full. If the filter is full, an item is removed to make room for the new
// item. This introduces a possibility for false negatives. To avoid this, use
// Count and Capacity to check if the filter is full before adding an item.
func (c *CuckooFilter) TestAndAdd(data []byte) (bool, error) {
	i1, i2, f := c.components(data)

	// If either bucket contains f, it's a member.
	if c.buckets[i1%c.m].contains(f) || c.buckets[i2%c.m].contains(f) {
		return true, nil
	}

	return false, c.add(i1, i2, f)
}

// TestAndRemove will test for membership of the data and remove it from the
// filter if it exists. Returns true if the data was a member, false if not.
func (c *CuckooFilter) TestAndRemove(data []byte) bool {
	i1, i2, f := c.components(data)

	// Try to remove from bucket[i1].
	b1 := c.buckets[i1%c.m]
	if idx := b1.indexOf(f); idx != -1 {
		b1[idx] = nil
		c.count--
		return true
	}

	// Try to remove from bucket[i2].
	b2 := c.buckets[i2%c.m]
	if idx := b2.indexOf(f); idx != -1 {
		b2[idx] = nil
		c.count--
		return true
	}

	return false
}

// Reset restores the Bloom filter to its original state. It returns the filter
// to allow for chaining.
func (c *CuckooFilter) Reset() *CuckooFilter {
	buckets := make([]bucket, c.m)
	for i := uint(0); i < c.m; i++ {
		buckets[i] = make(bucket, c.b)
	}
	c.buckets = buckets
	c.count = 0
	return c
}

// add will insert the fingerprint into the filter returning an error if the
// filter is full.
func (c *CuckooFilter) add(i1, i2 uint, f []byte) error {
	// Try to insert into bucket[i1].
	b1 := c.buckets[i1%c.m]
	if idx, err := b1.getEmptyEntry(); err == nil {
		b1[idx] = f
		c.count++
		return nil
	}

	// Try to insert into bucket[i2].
	b2 := c.buckets[i2%c.m]
	if idx, err := b2.getEmptyEntry(); err == nil {
		b2[idx] = f
		c.count++
		return nil
	}

	// Must relocate existing items.
	i := i1
	for n := 0; n < maxNumKicks; n++ {
		idx := i % c.m
		tmp := c.buckets[idx][rand.Intn(int(c.b))]
		c.buckets[idx][0] = f
		f = tmp
		i = i ^ uint(binary.BigEndian.Uint32(c.computeHash(f)))
		b := c.buckets[i%c.m]
		if idx, err := b.getEmptyEntry(); err == nil {
			b[idx] = f
			c.count++
			return nil
		}
	}

	return errors.New("full")
}

// components returns the two hash values used to index into the buckets and
// the fingerprint for the given element.
func (c *CuckooFilter) components(data []byte) (uint, uint, []byte) {
	var (
		hash = c.computeHash(data)
		f    = hash[0:c.f]
		i1   = uint(binary.BigEndian.Uint32(hash))
		i2   = i1 ^ uint(binary.BigEndian.Uint32(c.computeHash(f)))
	)

	return i1, i2, f
}

// computeHash returns a 32-bit hash value for the given data.
func (c *CuckooFilter) computeHash(data []byte) []byte {
	c.hash.Write(data)
	hash := c.hash.Sum(nil)
	c.hash.Reset()
	return hash
}

// calculateF returns the optimal fingerprint length in bytes for the given
// bucket size and false-positive rate epsilon.
func calculateF(b uint, epsilon float64) uint {
	f := uint(math.Ceil(math.Log(2 * float64(b) / epsilon)))
	f = f / 8
	if f <= 0 {
		f = 1
	}
	return f
}

// power2 calculates the next power of two for the given value.
func power2(x uint) uint {
	x--
	x |= x >> 1
	x |= x >> 2
	x |= x >> 4
	x |= x >> 8
	x |= x >> 16
	x |= x >> 32
	x++
	return x
}
<|endoftext|>