<commit_before>package main

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"github.com/rdwilliamson/aws/glacier"
	"io"
	"log"
	"os"
	"strconv"
	"time"
)

type retrievalData struct {
	Region       string
	Vault        string
	PartSize     uint64
	Job          string
	Downloaded   uint64
	Size         uint64
	FullTreeHash string
}

var (
	output string
	data   retrievalData
)

func (data *retrievalData) saveState(output string) {
	file, err := os.Create(output)
	if err != nil {
		log.Println("could not save state:", err)
		return
	}
	defer file.Close()
	enc := gob.NewEncoder(file)
	err = enc.Encode(data)
	if err != nil {
		log.Println("could not save state:", err)
		return
	}
}

func job(args []string) {
	if len(args) < 1 {
		fmt.Println("no job command")
		os.Exit(1)
	}
	command := args[0]
	args = args[1:]

	switch command {
	case "inventory":
		args = getConnection(args)

		if len(args) < 1 {
			fmt.Println("no vault")
			os.Exit(1)
		}
		vault := args[0]
		args = args[1:]

		var description, topic string
		if len(args) > 0 {
			topic = args[0]
		}
		if len(args) > 1 {
			description = args[1]
		}

		jobId, err := connection.InitiateInventoryJob(vault, topic, description)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		fmt.Println(jobId)

	case "archive":
		args = getConnection(args)

		if len(args) < 2 {
			fmt.Println("no vault")
			os.Exit(1)
		}
		vault := args[0]
		archive := args[1]
		args = args[2:]

		var description, topic string
		if len(args) > 0 {
			topic = args[0]
		}
		if len(args) > 1 {
			description = args[1]
		}

		jobId, err := connection.InitiateRetrievalJob(vault, archive, topic, description)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		fmt.Println(jobId)

	case "list":
		args = getConnection(args)

		if len(args) < 1 {
			fmt.Println("no vault")
			os.Exit(1)
		}
		vault := args[0]

		jobs, _, err := connection.ListJobs(vault, "", "", "", 0)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		for _, v := range jobs {
			fmt.Println("Action:", v.Action)
			if v.Action == "ArchiveRetrieval" {
				fmt.Println("Archive ID:", v.ArchiveId)
				fmt.Println("Archive Size:", v.ArchiveSizeInBytes, prettySize(v.ArchiveSizeInBytes))
			}
			fmt.Println("Completed:", v.Completed)
			if v.Completed {
				fmt.Println("Completion Date:", v.CompletionDate)
			}
			fmt.Println("Creation Date:", v.CreationDate)
			if v.Completed && v.Action == "InventoryRetrieval" {
				fmt.Println("Invenotry Size:", v.InventorySizeInBytes, prettySize(uint64(v.InventorySizeInBytes)))
			}
			fmt.Println("Job Description:", v.JobDescription)
			fmt.Println("Job ID:", v.JobId)
			if v.Action == "ArchiveRetrieval" {
				fmt.Println("SHA256 Tree Hash:", v.SHA256TreeHash)
			}
			fmt.Println("SNS Topic:", v.SNSTopic)
			fmt.Println("Status Code:", v.StatusCode)
			fmt.Println("Status Message:", v.StatusMessage)
			fmt.Println("Vault ARN:", v.VaultARN)
			fmt.Println()
		}

	case "describe":
		args = getConnection(args)

		if len(args) < 2 {
			fmt.Println("no vault and/or job id")
			os.Exit(1)
		}
		vault := args[0]
		jobId := args[1]

		job, err := connection.DescribeJob(vault, jobId)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		fmt.Println("Action:", job.Action)
		if job.Action == "ArchiveRetrieval" {
			fmt.Println("Archive ID:", job.ArchiveId)
			fmt.Println("Archive Size:", job.ArchiveSizeInBytes, prettySize(job.ArchiveSizeInBytes))
		}
		fmt.Println("Completed:", job.Completed)
		if job.Completed {
			fmt.Println("Completion Date:", job.CompletionDate)
		}
		fmt.Println("Creation Date:", job.CreationDate)
		if job.Completed && job.Action == "InventoryRetrieval" {
			fmt.Println("Invenotry Size:", job.InventorySizeInBytes, prettySize(uint64(job.InventorySizeInBytes)))
		}
		fmt.Println("Job Description:", job.JobDescription)
		fmt.Println("Job ID:", job.JobId)
		if job.Action == "ArchiveRetrieval" {
			fmt.Println("SHA256 Tree Hash:", job.SHA256TreeHash)
		}
		fmt.Println("SNS Topic:", job.SNSTopic)
		fmt.Println("Status Code:", job.StatusCode)
		fmt.Println("Status Message:", job.StatusMessage)
		fmt.Println("Vault ARN:", job.VaultARN)

	case "get":
		if len(args) < 1 {
			fmt.Println("no job sub command")
			os.Exit(1)
		}
		subCommand := args[0]
		args = args[1:]

		switch subCommand {
		case "inventory":
			args = getConnection(args)

			if len(args) < 2 {
				fmt.Println("no vault and/or job id")
				os.Exit(1)
			}
			vault := args[0]
			job := args[1]

			inventory, err := connection.GetInventoryJob(vault, job)
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}

			fmt.Println("Vault ARN:", inventory.VaultARN)
			fmt.Println("Inventory Date:", inventory.InventoryDate)

			for _, v := range inventory.ArchiveList {
				fmt.Println()
				fmt.Println("Archive ID:", v.ArchiveId)
				fmt.Println("Archive Description:", v.ArchiveDescription)
				fmt.Println("Creation Date:", v.CreationDate)
				fmt.Println("Size:", v.Size, prettySize(v.Size))
				fmt.Println("SHA256 Tree Hash:", v.SHA256TreeHash)
			}

		case "archive":
			args = getConnection(args)

			if len(args) < 3 {
				fmt.Println("no vault, job id, and/or output file")
				os.Exit(1)
			}
			vault := args[0]
			job := args[1]
			fileName := args[2]

			file, err := os.Create(fileName)
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			defer file.Close()

			archive, _, err := connection.GetRetrievalJob(vault, job, 0, 0)
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			defer archive.Close()

			_, err = io.Copy(file, archive)
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
		default:
			fmt.Println("unknown job sub command:", subCommand)
			os.Exit(1)
		}

	case "run":
		args = getConnection(args)
		if len(args) < 4 {
			fmt.Println("no vault, archive, download size and/or output file")
			os.Exit(1)
		}
		vault := args[0]
		archive := args[1]
		partSize, err := strconv.ParseUint(args[2], 10, 64)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		partSize *= 1024 * 1024
		output = args[3]
		args = args[4:]

		var topic string
		if len(args) > 0 {
			topic = args[0]
			args = args[1:]
		}

		var description string
		if len(args) > 0 {
			description = args[0]
			args = args[1:]
		}

		// initiate retrieval job
		job, err := connection.InitiateRetrievalJob(vault, archive, topic, description)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		log.Println("initiated retrieval job:", job)

		// save state
		data.Region = connection.Signature.Region.Name
		data.Vault = vault
		data.PartSize = partSize
		data.Job = job
		data.saveState(output + ".gob")

		// wait for job to complete, using polling
		time.Sleep(3 * time.Hour)

		// check status sleeping 15m?
		var try int
		for {
			job, err := connection.DescribeJob(vault, job)
			if err != nil {
				log.Println(err)
				try++
				if try > retries {
					fmt.Println("too many retries")
					os.Exit(1)
				}
			} else {
				try = 0
				if job.Completed {
					data.Size = uint64(job.ArchiveSizeInBytes)
					data.FullTreeHash = job.SHA256TreeHash
					data.saveState(output + ".gob")
					break
				}
				log.Println("retrieval job not yet completed")
				time.Sleep(15 * time.Minute)
			}
		}

		fallthrough
	case "resume":
		if command == "resume" {
			if len(args) < 1 {
				fmt.Println("no filename")
				os.Exit(1)
			}
			output = args[0]
			args = args[1:]

			file, err := os.Open(output + ".gob")
			if err != nil {
				fmt.Println("could not resume:", err)
				os.Exit(1)
			}
			dec := gob.NewDecoder(file)
			err = dec.Decode(&data)
			file.Close()
			if err != nil {
				fmt.Println("could not resume:", err)
				os.Exit(1)
			}

			getConnection([]string{data.Region})

			if len(args) > 0 {
				data.PartSize, err = strconv.ParseUint(args[0], 10, 64)
				if err != nil {
					fmt.Println(err)
					os.Exit(1)
				}
				data.PartSize *= 1024 * 1024
				data.saveState(output + ".gob")
			}
		}

		file, err := os.OpenFile(output, os.O_RDWR|os.O_APPEND|os.O_CREATE, 0666)
		if err != nil {
			log.Println(err)
			os.Exit(1)
		}
		defer file.Close()

		// loop getting parts, checking tree hash of each
		buffer := bytes.NewBuffer(make([]byte, data.PartSize))
		hasher := glacier.NewTreeHash()
		var try int

		if command == "resume" {
			_, err = file.Seek(int64(data.Downloaded), 0)
			if err != nil {
				fmt.Println("could not resume:", err)
				os.Exit(1)
			}
		}

		for data.Downloaded < data.Size {
			log.Println("downloading", data.Downloaded, "to", data.Downloaded+data.PartSize-1, "of", data.Size)
			buffer.Reset()
			hasher.Reset()

			part, treeHash, err := connection.GetRetrievalJob(data.Vault, data.Job, data.Downloaded,
				data.Downloaded+data.PartSize-1)
			if err != nil {
				log.Println("GetRetrievalJob:", err)
				try++
				if try > retries {
					fmt.Println("too many retries")
					os.Exit(1)
				}
				continue
			}

			// copy to temporary buffer
			_, err = io.Copy(buffer, part)
			if err != nil {
				log.Println("io.Copy:", err)
				try++
				if try > retries {
					fmt.Println("too many retries")
					os.Exit(1)
				}
				continue
			}

			// check tree hash
			// TODO only if size is MiB power of two and partial content aligns
			// on a MiB
			_, err = io.CopyN(hasher, buffer, int64(data.PartSize))
			if err != nil && err != io.EOF {
				log.Println("hashing", err)
				try++
				if try > retries {
					fmt.Println("too many retries")
					os.Exit(1)
				}
				continue
			}
			hasher.Close()
			if treeHash != hasher.TreeHash() {
				log.Println("tree hash mismatch, want", treeHash, "got", hasher.TreeHash())
				try++
				if try > retries {
					fmt.Println("too many retries")
					os.Exit(1)
				}
				continue
			}
			log.Println("checked tree hash")

			// copy to file
			_, err = file.Write(buffer.Bytes())
			if err != nil {
				log.Println("copying buffer to file:", err)
				try++
				if try > retries {
					fmt.Println("too many retries")
					os.Exit(1)
				}
			}
			log.Println("copied to file")

			// save state
			data.Downloaded += data.PartSize
			data.saveState(output + ".gob")

			try = 0
		}

		// check tree hash of entire archive
		log.Println("download complete, verifying")
		_, err = file.Seek(0, 0)
		if err != nil {
			log.Println("seek:", err)
			os.Exit(1)
		}

		hasher.Reset()
		_, err = io.Copy(hasher, file)
		if err != nil {
			log.Println("hashing whole file:", err)
			os.Exit(1)
		}
		hasher.Close()

		if hasher.TreeHash() != data.FullTreeHash {
			log.Println("entire file tree hash mismatch, want", data.FullTreeHash, "got", hasher.TreeHash())
			os.Exit(1)
		}

		os.Remove(output + ".gob")

	default:
		fmt.Println("unknown job command:", command)
		os.Exit(1)
	}
}
<commit_msg>Avoid treeHash errors when we don't have anything to compare to.<commit_after>package main

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"github.com/rdwilliamson/aws/glacier"
	"io"
	"log"
	"os"
	"strconv"
	"time"
)

type retrievalData struct {
	Region       string
	Vault        string
	PartSize     uint64
	Job          string
	Downloaded   uint64
	Size         uint64
	FullTreeHash string
}

var (
	output string
	data   retrievalData
)

func (data *retrievalData) saveState(output string) {
	file, err := os.Create(output)
	if err != nil {
		log.Println("could not save state:", err)
		return
	}
	defer file.Close()
	enc := gob.NewEncoder(file)
	err = enc.Encode(data)
	if err != nil {
		log.Println("could not save state:", err)
		return
	}
}

func job(args []string) {
	if len(args) < 1 {
		fmt.Println("no job command")
		os.Exit(1)
	}
	command := args[0]
	args = args[1:]

	switch command {
	case "inventory":
		args = getConnection(args)

		if len(args) < 1 {
			fmt.Println("no vault")
			os.Exit(1)
		}
		vault := args[0]
		args = args[1:]

		var description, topic string
		if len(args) > 0 {
			topic = args[0]
		}
		if len(args) > 1 {
			description = args[1]
		}

		jobId, err := connection.InitiateInventoryJob(vault, topic, description)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		fmt.Println(jobId)

	case "archive":
		args = getConnection(args)

		if len(args) < 2 {
			fmt.Println("no vault")
			os.Exit(1)
		}
		vault := args[0]
		archive := args[1]
		args = args[2:]

		var description, topic string
		if len(args) > 0 {
			topic = args[0]
		}
		if len(args) > 1 {
			description = args[1]
		}

		jobId, err := connection.InitiateRetrievalJob(vault, archive, topic, description)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		fmt.Println(jobId)

	case "list":
		args = getConnection(args)

		if len(args) < 1 {
			fmt.Println("no vault")
			os.Exit(1)
		}
		vault := args[0]

		jobs, _, err := connection.ListJobs(vault, "", "", "", 0)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		for _, v := range jobs {
			fmt.Println("Action:", v.Action)
			if v.Action == "ArchiveRetrieval" {
				fmt.Println("Archive ID:", v.ArchiveId)
				fmt.Println("Archive Size:", v.ArchiveSizeInBytes, prettySize(v.ArchiveSizeInBytes))
			}
			fmt.Println("Completed:", v.Completed)
			if v.Completed {
				fmt.Println("Completion Date:", v.CompletionDate)
			}
			fmt.Println("Creation Date:", v.CreationDate)
			if v.Completed && v.Action == "InventoryRetrieval" {
				fmt.Println("Invenotry Size:", v.InventorySizeInBytes, prettySize(uint64(v.InventorySizeInBytes)))
			}
			fmt.Println("Job Description:", v.JobDescription)
			fmt.Println("Job ID:", v.JobId)
			if v.Action == "ArchiveRetrieval" {
				fmt.Println("SHA256 Tree Hash:", v.SHA256TreeHash)
			}
			fmt.Println("SNS Topic:", v.SNSTopic)
			fmt.Println("Status Code:", v.StatusCode)
			fmt.Println("Status Message:", v.StatusMessage)
			fmt.Println("Vault ARN:", v.VaultARN)
			fmt.Println()
		}

	case "describe":
		args = getConnection(args)

		if len(args) < 2 {
			fmt.Println("no vault and/or job id")
			os.Exit(1)
		}
		vault := args[0]
		jobId := args[1]

		job, err := connection.DescribeJob(vault, jobId)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		fmt.Println("Action:", job.Action)
		if job.Action == "ArchiveRetrieval" {
			fmt.Println("Archive ID:", job.ArchiveId)
			fmt.Println("Archive Size:", job.ArchiveSizeInBytes, prettySize(job.ArchiveSizeInBytes))
		}
		fmt.Println("Completed:", job.Completed)
		if job.Completed {
			fmt.Println("Completion Date:", job.CompletionDate)
		}
		fmt.Println("Creation Date:", job.CreationDate)
		if job.Completed && job.Action == "InventoryRetrieval" {
			fmt.Println("Invenotry Size:", job.InventorySizeInBytes, prettySize(uint64(job.InventorySizeInBytes)))
		}
		fmt.Println("Job Description:", job.JobDescription)
		fmt.Println("Job ID:", job.JobId)
		if job.Action == "ArchiveRetrieval" {
			fmt.Println("SHA256 Tree Hash:", job.SHA256TreeHash)
		}
		fmt.Println("SNS Topic:", job.SNSTopic)
		fmt.Println("Status Code:", job.StatusCode)
		fmt.Println("Status Message:", job.StatusMessage)
		fmt.Println("Vault ARN:", job.VaultARN)

	case "get":
		if len(args) < 1 {
			fmt.Println("no job sub command")
			os.Exit(1)
		}
		subCommand := args[0]
		args = args[1:]

		switch subCommand {
		case "inventory":
			args = getConnection(args)

			if len(args) < 2 {
				fmt.Println("no vault and/or job id")
				os.Exit(1)
			}
			vault := args[0]
			job := args[1]

			inventory, err := connection.GetInventoryJob(vault, job)
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}

			fmt.Println("Vault ARN:", inventory.VaultARN)
			fmt.Println("Inventory Date:", inventory.InventoryDate)

			for _, v := range inventory.ArchiveList {
				fmt.Println()
				fmt.Println("Archive ID:", v.ArchiveId)
				fmt.Println("Archive Description:", v.ArchiveDescription)
				fmt.Println("Creation Date:", v.CreationDate)
				fmt.Println("Size:", v.Size, prettySize(v.Size))
				fmt.Println("SHA256 Tree Hash:", v.SHA256TreeHash)
			}

		case "archive":
			args = getConnection(args)

			if len(args) < 3 {
				fmt.Println("no vault, job id, and/or output file")
				os.Exit(1)
			}
			vault := args[0]
			job := args[1]
			fileName := args[2]

			file, err := os.Create(fileName)
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			defer file.Close()

			archive, _, err := connection.GetRetrievalJob(vault, job, 0, 0)
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			defer archive.Close()

			_, err = io.Copy(file, archive)
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
		default:
			fmt.Println("unknown job sub command:", subCommand)
			os.Exit(1)
		}

	case "run":
		args = getConnection(args)
		if len(args) < 4 {
			fmt.Println("no vault, archive, download size and/or output file")
			os.Exit(1)
		}
		vault := args[0]
		archive := args[1]
		partSize, err := strconv.ParseUint(args[2], 10, 64)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		partSize *= 1024 * 1024
		output = args[3]
		args = args[4:]

		var topic string
		if len(args) > 0 {
			topic = args[0]
			args = args[1:]
		}

		var description string
		if len(args) > 0 {
			description = args[0]
			args = args[1:]
		}

		// initiate retrieval job
		job, err := connection.InitiateRetrievalJob(vault, archive, topic, description)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		log.Println("initiated retrieval job:", job)

		// save state
		data.Region = connection.Signature.Region.Name
		data.Vault = vault
		data.PartSize = partSize
		data.Job = job
		data.saveState(output + ".gob")

		// wait for job to complete, using polling
		time.Sleep(3 * time.Hour)

		// check status sleeping 15m?
		var try int
		for {
			job, err := connection.DescribeJob(vault, job)
			if err != nil {
				log.Println(err)
				try++
				if try > retries {
					fmt.Println("too many retries")
					os.Exit(1)
				}
			} else {
				try = 0
				if job.Completed {
					data.Size = uint64(job.ArchiveSizeInBytes)
					data.FullTreeHash = job.SHA256TreeHash
					data.saveState(output + ".gob")
					break
				}
				log.Println("retrieval job not yet completed")
				time.Sleep(15 * time.Minute)
			}
		}

		fallthrough
	case "resume":
		if command == "resume" {
			if len(args) < 1 {
				fmt.Println("no filename")
				os.Exit(1)
			}
			output = args[0]
			args = args[1:]

			file, err := os.Open(output + ".gob")
			if err != nil {
				fmt.Println("could not resume:", err)
				os.Exit(1)
			}
			dec := gob.NewDecoder(file)
			err = dec.Decode(&data)
			file.Close()
			if err != nil {
				fmt.Println("could not resume:", err)
				os.Exit(1)
			}

			getConnection([]string{data.Region})

			if len(args) > 0 {
				data.PartSize, err = strconv.ParseUint(args[0], 10, 64)
				if err != nil {
					fmt.Println(err)
					os.Exit(1)
				}
				data.PartSize *= 1024 * 1024
				data.saveState(output + ".gob")
			}
		}

		file, err := os.OpenFile(output, os.O_RDWR|os.O_APPEND|os.O_CREATE, 0666)
		if err != nil {
			log.Println(err)
			os.Exit(1)
		}
		defer file.Close()

		// loop getting parts, checking tree hash of each
		buffer := bytes.NewBuffer(make([]byte, data.PartSize))
		hasher := glacier.NewTreeHash()
		var try int

		if command == "resume" {
			_, err = file.Seek(int64(data.Downloaded), 0)
			if err != nil {
				fmt.Println("could not resume:", err)
				os.Exit(1)
			}
		}

		for data.Downloaded < data.Size {
			log.Println("downloading", data.Downloaded, "to", data.Downloaded+data.PartSize-1, "of", data.Size)
			buffer.Reset()
			hasher.Reset()

			part, treeHash, err := connection.GetRetrievalJob(data.Vault, data.Job, data.Downloaded,
				data.Downloaded+data.PartSize-1)
			if err != nil {
				log.Println("GetRetrievalJob:", err)
				try++
				if try > retries {
					fmt.Println("too many retries")
					os.Exit(1)
				}
				continue
			}

			// copy to temporary buffer
			_, err = io.Copy(buffer, part)
			if err != nil {
				log.Println("io.Copy:", err)
				try++
				if try > retries {
					fmt.Println("too many retries")
					os.Exit(1)
				}
				continue
			}

			// check tree hash
			_, err = io.CopyN(hasher, buffer, int64(data.PartSize))
			if err != nil && err != io.EOF {
				log.Println("hashing", err)
				try++
				if try > retries {
					fmt.Println("too many retries")
					os.Exit(1)
				}
				continue
			}
			hasher.Close()
			if treeHash != "" && treeHash != hasher.TreeHash() {
				log.Println("tree hash mismatch, want", treeHash, "got", hasher.TreeHash())
				try++
				if try > retries {
					fmt.Println("too many retries")
					os.Exit(1)
				}
				continue
			}
			log.Println("checked tree hash")

			// copy to file
			_, err = file.Write(buffer.Bytes())
			if err != nil {
				log.Println("copying buffer to file:", err)
				try++
				if try > retries {
					fmt.Println("too many retries")
					os.Exit(1)
				}
			}
			log.Println("copied to file")

			// save state
			data.Downloaded += data.PartSize
			data.saveState(output + ".gob")

			try = 0
		}

		// check tree hash of entire archive
		log.Println("download complete, verifying")
		_, err = file.Seek(0, 0)
		if err != nil {
			log.Println("seek:", err)
			os.Exit(1)
		}

		hasher.Reset()
		_, err = io.Copy(hasher, file)
		if err != nil {
			log.Println("hashing whole file:", err)
			os.Exit(1)
		}
		hasher.Close()

		if hasher.TreeHash() != data.FullTreeHash {
			log.Println("entire file tree hash mismatch, want", data.FullTreeHash, "got", hasher.TreeHash())
			os.Exit(1)
		}

		os.Remove(output + ".gob")

	default:
		fmt.Println("unknown job command:", command)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2015 Thomas de Zeeuw.
//
// Licensed under the MIT license that can be found in the LICENSE file.

package logger

import (
	"fmt"
	"runtime"
	"sync"
	"time"

	"github.com/Thomasdezeeuw/logger/internal/util"
)

const (
	defaultStackSize        = 8192
	defaultEventChannelSize = 1024
	maxNWriteErrors         = 5
)

// EventWriter is the backend of the logger package. It takes events and writes
// them to a storage, this can for example be a file, a database or the
// terminal.
type EventWriter interface {
	// Write is called every time a Log operation is called. This call should
	// block until the event is written. Advised is to buffer the events. All
	// calls to Write are synchronous.
	//
	// If an error is returned the event is expected to NOT have been written. The
	// event will written again after the error handler is called. If the
	// EventWriter returns 5 errors in a row the EventWriter is consided to be bad
	// and will be removed from the list of event writers. HandlerError will be
	// called with ErrBadEventWriter if this happens.
	Write(Event) error

	// HandleError is called every time Write returns an error.
	HandleError(error)

	// Close is called on the EventWriter once Close() (on the package) is called.
	Close() error
}

var (
	eventChannel       = make(chan Event, defaultEventChannelSize)
	eventChannelClosed = make(chan struct{}, 1) // Can't block.
	eventWriters       []EventWriter
	started            bool
)

// Start starts the logger package and enables writing to the given
// EventWriters.
func Start(ews ...EventWriter) {
	if started {
		panic("logger: can only Start once")
	} else if len(ews) < 1 {
		panic("logger: need atleast a single EventWriter to write to")
	}

	started = true
	eventWriters = ews

	go writeEvents()
}

var ErrBadEventWriter = fmt.Errorf("EventWriter is bad, %d faulty writes, EventWriter will be dropped", maxNWriteErrors)

// Needs to be run in it's own goroutine, it blocks until eventChannel is
// closed. After eventChannel is closed it sends a signal to eventChannelClosed.
func writeEvents() {
	var wg sync.WaitGroup
	wg.Add(len(eventWriters))

	// Create event sub channels for each EventWriter and start each EventWriter.
	var eventSubChannels = make([]chan Event, len(eventWriters))
	for i, ew := range eventWriters {
		eventSubChannels[i] = make(chan Event, defaultEventChannelSize)
		go startEventWriter(ew, eventSubChannels[i], &wg)
	}

	// Fan out the events to all the sub channels.
	for event := range eventChannel {
		for _, eventSubChannel := range eventSubChannels {
			eventSubChannel <- event
		}
	}

	// Close each sub channel.
	for _, eventSubChannel := range eventSubChannels {
		close(eventSubChannel)
	}

	wg.Wait()
	eventChannelClosed <- struct{}{}
}

// StartEventWriter blocks until the events channel is closed.
func startEventWriter(ew EventWriter, events <-chan Event, wg *sync.WaitGroup) {
	var badEventWriter = false

	for event := range events {
		if badEventWriter {
			// Simply drain the channel.
			// todo: improve this, don't send to the channel anymore if the writer is
			// bad.
			continue
		}

		if err := writeEvent(ew, event); err != nil {
			// At this point the EventWriter is bad and we won't write to it anymore.
			ew.HandleError(err)
			badEventWriter = true
		}
	}

	wg.Done()
}

// WriteEvent tries to write the event to the given EventWriter, it tries it up
// to maxNWriteErrors times. If EventWriter.Write returns an error it gets
// passed to the error handler of the EventWriter. This function either returns
// ErrBadEventWriter or nil as an error.
func writeEvent(ew EventWriter, event Event) error {
	for n := 1; n <= maxNWriteErrors; n++ {
		err := ew.Write(event)
		if err == nil {
			return nil
		}

		// Handle the error and try again.
		ew.HandleError(err)
	}

	return ErrBadEventWriter
}

// Close stops all the Log Operations from being usable, and they will panic if
// used after Close is called. It also closes all EventWriters and returns the
// first returned error. The EventWriters are closed in the order they are
// passed to Start.
func Close() error {
	close(eventChannel)
	<-eventChannelClosed

	var err error
	for _, eventWriter := range eventWriters {
		er := eventWriter.Close()
		if er != nil && err == nil {
			err = er
		}
	}
	return err
}

// Subbed for testing.
var now = func() time.Time {
	return time.Now()
}

// Debug logs a debug message.
func Debug(tags Tags, msg string) {
	eventChannel <- Event{DebugEvent, now(), tags, msg, nil}
}

// Debugf is a formatted function of Debug.
func Debugf(tags Tags, format string, v ...interface{}) {
	Debug(tags, fmt.Sprintf(format, v...))
}

// Info logs an informational message.
func Info(tags Tags, msg string) {
	eventChannel <- Event{InfoEvent, now(), tags, msg, nil}
}

// Infof is a formatted function of Info.
func Infof(tags Tags, format string, v ...interface{}) {
	Info(tags, fmt.Sprintf(format, v...))
}

// Warn logs a warning message.
func Warn(tags Tags, msg string) {
	eventChannel <- Event{WarnEvent, now(), tags, msg, nil}
}

// Warnf is a formatted function of Warn.
func Warnf(tags Tags, format string, v ...interface{}) {
	Warn(tags, fmt.Sprintf(format, v...))
}

// Error logs an error message.
func Error(tags Tags, err error) {
	eventChannel <- Event{ErrorEvent, now(), tags, err.Error(), nil}
}

// Errorf is a formatted function of Error.
func Errorf(tags Tags, format string, v ...interface{}) {
	Error(tags, fmt.Errorf(format, v...))
}

// Fatal logs a recovered error which could have killed the application. Fatal
// adds a stack trace (type []byte) as Event.Data.
func Fatal(tags Tags, recv interface{}) {
	stackTrace := make([]byte, defaultStackSize)
	n := runtime.Stack(stackTrace, false)
	stackTrace = stackTrace[:n]

	msg := util.InterfaceToString(recv)
	eventChannel <- Event{FatalEvent, now(), tags, msg, stackTrace}
}

// Thumbstone indicates a function is still used in production. When developing
// software it's possible to introduce dead code with updates and new features.
// If a function is being suspected of being dead (not used) in production, add
// a call to Thumbstone and check the production logs to see if you're right.
//
// The caller of the (possibly) dead function will be put in the message, using
// the following format:
//	Function functionName called by callerFunctionName, from file /path/to/file on line lineNumber
// For example:
//	Function myFunction called by main.main, from file /main.go on line 20
func Thumbstone(tags Tags, functionName string) {
	var msg string
	if pc, file, line, ok := runtime.Caller(2); ok {
		fn := runtime.FuncForPC(pc)
		msg = fmt.Sprintf("Function %s called by %s, from file %s on line %d",
			functionName, fn.Name(), file, line)
	} else {
		msg = "Function " + functionName + " called from unkown location"
	}

	eventChannel <- Event{ThumbEvent, now(), tags, msg, nil}
}

// Log logs a custom created event.
//
// Note: the timestamp doesn't need to be set, because it will be set by Log.
func Log(event Event) {
	event.Timestamp = now()
	eventChannel <- event
}
<commit_msg>Doc ErrBadEventWriter<commit_after>// Copyright (C) 2015 Thomas de Zeeuw.
//
// Licensed under the MIT license that can be found in the LICENSE file.

package logger

import (
	"fmt"
	"runtime"
	"sync"
	"time"

	"github.com/Thomasdezeeuw/logger/internal/util"
)

const (
	defaultStackSize        = 8192
	defaultEventChannelSize = 1024
	maxNWriteErrors         = 5
)

// EventWriter is the backend of the logger package. It takes events and writes
// them to a storage, this can for example be a file, a database or the
// terminal.
type EventWriter interface {
	// Write is called every time a Log operation is called. This call should
	// block until the event is written. Advised is to buffer the events. All
	// calls to Write are synchronous.
	//
	// If an error is returned the event is expected to NOT have been written. The
	// event will written again after the error handler is called. If the
	// EventWriter returns 5 errors in a row the EventWriter is consided to be bad
	// and will be removed from the list of event writers. HandlerError will be
	// called with ErrBadEventWriter if this happens.
	Write(Event) error

	// HandleError is called every time Write returns an error.
	HandleError(error)

	// Close is called on the EventWriter once Close() (on the package) is called.
	Close() error
}

var (
	eventChannel       = make(chan Event, defaultEventChannelSize)
	eventChannelClosed = make(chan struct{}, 1) // Can't block.
	eventWriters       []EventWriter
	started            bool
)

// Start starts the logger package and enables writing to the given
// EventWriters.
func Start(ews ...EventWriter) {
	if started {
		panic("logger: can only Start once")
	} else if len(ews) < 1 {
		panic("logger: need atleast a single EventWriter to write to")
	}

	started = true
	eventWriters = ews

	go writeEvents()
}

// ErrBadEventWriter gets passed to the error handler of an EventWriter after it
// returned too many write errors in a row. After the error handler of the
// EventWriter is called with this error the writer is considered faulty and
// will no longer recive any Events.
var ErrBadEventWriter = fmt.Errorf("EventWriter is bad, %d faulty writes, EventWriter will be dropped", maxNWriteErrors)

// Needs to be run in it's own goroutine, it blocks until eventChannel is
// closed. After eventChannel is closed it sends a signal to eventChannelClosed.
func writeEvents() {
	var wg sync.WaitGroup
	wg.Add(len(eventWriters))

	// Create event sub channels for each EventWriter and start each EventWriter.
	var eventSubChannels = make([]chan Event, len(eventWriters))
	for i, ew := range eventWriters {
		eventSubChannels[i] = make(chan Event, defaultEventChannelSize)
		go startEventWriter(ew, eventSubChannels[i], &wg)
	}

	// Fan out the events to all the sub channels.
	for event := range eventChannel {
		for _, eventSubChannel := range eventSubChannels {
			eventSubChannel <- event
		}
	}

	// Close each sub channel.
	for _, eventSubChannel := range eventSubChannels {
		close(eventSubChannel)
	}

	wg.Wait()
	eventChannelClosed <- struct{}{}
}

// StartEventWriter blocks until the events channel is closed.
func startEventWriter(ew EventWriter, events <-chan Event, wg *sync.WaitGroup) {
	var badEventWriter = false

	for event := range events {
		if badEventWriter {
			// Simply drain the channel.
			// todo: improve this, don't send to the channel anymore if the writer is
			// bad.
			continue
		}

		if err := writeEvent(ew, event); err != nil {
			// At this point the EventWriter is bad and we won't write to it anymore.
			ew.HandleError(err)
			badEventWriter = true
		}
	}

	wg.Done()
}

// WriteEvent tries to write the event to the given EventWriter, it tries it up
// to maxNWriteErrors times. If EventWriter.Write returns an error it gets
// passed to the error handler of the EventWriter. This function either returns
// ErrBadEventWriter or nil as an error.
func writeEvent(ew EventWriter, event Event) error {
	for n := 1; n <= maxNWriteErrors; n++ {
		err := ew.Write(event)
		if err == nil {
			return nil
		}

		// Handle the error and try again.
		ew.HandleError(err)
	}

	return ErrBadEventWriter
}

// Close stops all the Log Operations from being usable, and they will panic if
// used after Close is called. It also closes all EventWriters and returns the
// first returned error. The EventWriters are closed in the order they are
// passed to Start.
func Close() error {
	close(eventChannel)
	<-eventChannelClosed

	var err error
	for _, eventWriter := range eventWriters {
		er := eventWriter.Close()
		if er != nil && err == nil {
			err = er
		}
	}
	return err
}

// Subbed for testing.
var now = func() time.Time {
	return time.Now()
}

// Debug logs a debug message.
func Debug(tags Tags, msg string) {
	eventChannel <- Event{DebugEvent, now(), tags, msg, nil}
}

// Debugf is a formatted function of Debug.
func Debugf(tags Tags, format string, v ...interface{}) {
	Debug(tags, fmt.Sprintf(format, v...))
}

// Info logs an informational message.
func Info(tags Tags, msg string) {
	eventChannel <- Event{InfoEvent, now(), tags, msg, nil}
}

// Infof is a formatted function of Info.
func Infof(tags Tags, format string, v ...interface{}) {
	Info(tags, fmt.Sprintf(format, v...))
}

// Warn logs a warning message.
func Warn(tags Tags, msg string) {
	eventChannel <- Event{WarnEvent, now(), tags, msg, nil}
}

// Warnf is a formatted function of Warn.
func Warnf(tags Tags, format string, v ...interface{}) {
	Warn(tags, fmt.Sprintf(format, v...))
}

// Error logs an error message.
func Error(tags Tags, err error) {
	eventChannel <- Event{ErrorEvent, now(), tags, err.Error(), nil}
}

// Errorf is a formatted function of Error.
func Errorf(tags Tags, format string, v ...interface{}) {
	Error(tags, fmt.Errorf(format, v...))
}

// Fatal logs a recovered error which could have killed the application. Fatal
// adds a stack trace (type []byte) as Event.Data.
func Fatal(tags Tags, recv interface{}) {
	stackTrace := make([]byte, defaultStackSize)
	n := runtime.Stack(stackTrace, false)
	stackTrace = stackTrace[:n]

	msg := util.InterfaceToString(recv)
	eventChannel <- Event{FatalEvent, now(), tags, msg, stackTrace}
}

// Thumbstone indicates a function is still used in production. When developing
// software it's possible to introduce dead code with updates and new features.
// If a function is being suspected of being dead (not used) in production, add
// a call to Thumbstone and check the production logs to see if you're right.
//
// The caller of the (possibly) dead function will be put in the message, using
// the following format:
//	Function functionName called by callerFunctionName, from file /path/to/file on line lineNumber
// For example:
//	Function myFunction called by main.main, from file /main.go on line 20
func Thumbstone(tags Tags, functionName string) {
	var msg string
	if pc, file, line, ok := runtime.Caller(2); ok {
		fn := runtime.FuncForPC(pc)
		msg = fmt.Sprintf("Function %s called by %s, from file %s on line %d",
			functionName, fn.Name(), file, line)
	} else {
		msg = "Function " + functionName + " called from unkown location"
	}

	eventChannel <- Event{ThumbEvent, now(), tags, msg, nil}
}

// Log logs a custom created event.
//
// Note: the timestamp doesn't need to be set, because it will be set by Log.
func Log(event Event) {
	event.Timestamp = now()
	eventChannel <- event
}
<|endoftext|>


<commit_before>package buffer

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"io"
)

type Memory struct {
	N int64
	*bytes.Buffer
}

func New(n int64) *Memory {
	return &Memory{
		N:      n,
		Buffer: bytes.NewBuffer(nil),
	}
}

func (buf *Memory) Cap() int64 {
	return buf.N
}

func (buf *Memory) Len() int64 {
	return int64(buf.Buffer.Len())
}

func (buf *Memory) Write(p []byte) (n int, err error) {
	return LimitWriter(buf.Buffer, Gap(buf)).Write(p)
}

func (buf *Memory) WriteTo(w io.Writer) (n int64, err error) {
	return buf.Buffer.WriteTo(w)
}

// Must Add io.ErrShortWrite
func (buf *Memory) WriteAt(p []byte, off int64) (n int, err error) {
	if off > buf.Len() {
		return 0, bytes.ErrTooLarge
	} else if len64(p)+off <= buf.Len() {
		d := buf.Bytes()[off:]
		return copy(d, p), nil
	} else {
		d := buf.Bytes()[off:]
		n = copy(d, p)
		m, err := buf.Write(p[n:])
		return n + m, err
	}
}

func (buf *Memory) ReadAt(p []byte, off int64) (n int, err error) {
	return bytes.NewReader(buf.Bytes()).ReadAt(p, off)
}

func (buf *Memory) Read(p []byte) (n int, err error) {
	return io.LimitReader(buf.Buffer, buf.Len()).Read(p)
}

func (buf *Memory) ReadFrom(r io.Reader) (n int64, err error) {
	return buf.Buffer.ReadFrom(io.LimitReader(r, Gap(buf)))
}

func (buf *Memory) FFwd(n int64) int64 {
	data := buf.Bytes()
	if n > len64(data) {
		n = len64(data)
	}
	b := bytes.NewBuffer(data[n:])
	buf.Buffer = b
	return n
}

func init() {
	gob.Register(&Memory{})
}

func (buf *Memory) MarshalBinary() ([]byte, error) {
	var b bytes.Buffer
	fmt.Fprint(&b, buf.N)
	b.Write(buf.Bytes())
	return b.Bytes(), nil
}

func (buf *Memory) UnmarshalBinary(data []byte) error {
	b := bytes.NewBuffer(data)
	_, err := fmt.Fscan(b, &buf.N)
	buf.Buffer = b
	return err
}
<commit_msg>WriteTo was already supported by embeded bytes.Buffer<commit_after>package buffer

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"io"
)

type Memory struct {
	N int64
	*bytes.Buffer
}

func New(n int64) *Memory {
	return &Memory{
		N:      n,
		Buffer: bytes.NewBuffer(nil),
	}
}

func (buf *Memory) Cap() int64 {
	return buf.N
}

func (buf *Memory) Len() int64 {
	return int64(buf.Buffer.Len())
}

func (buf *Memory) Write(p []byte) (n int, err error) {
	return LimitWriter(buf.Buffer, Gap(buf)).Write(p)
}

// Must Add io.ErrShortWrite
func (buf *Memory) WriteAt(p []byte, off int64) (n int, err error) {
	if off > buf.Len() {
		return 0, bytes.ErrTooLarge
	} else if len64(p)+off <= buf.Len() {
		d := buf.Bytes()[off:]
		return copy(d, p), nil
	} else {
		d := buf.Bytes()[off:]
		n = copy(d, p)
		m, err := buf.Write(p[n:])
		return n + m, err
	}
}

func (buf *Memory) ReadAt(p []byte, off int64) (n int, err error) {
	return bytes.NewReader(buf.Bytes()).ReadAt(p, off)
}

func (buf *Memory) Read(p []byte) (n int, err error) {
	return io.LimitReader(buf.Buffer, buf.Len()).Read(p)
}

func (buf *Memory) ReadFrom(r io.Reader) (n int64, err error) {
	return buf.Buffer.ReadFrom(io.LimitReader(r, Gap(buf)))
}

func (buf *Memory) FFwd(n int64) int64 {
	data := buf.Bytes()
	if n > len64(data) {
		n = len64(data)
	}
	b := bytes.NewBuffer(data[n:])
	buf.Buffer = b
	return n
}

func init() {
	gob.Register(&Memory{})
}

func (buf *Memory) MarshalBinary() ([]byte, error) {
	var b bytes.Buffer
	fmt.Fprint(&b, buf.N)
	b.Write(buf.Bytes())
	return b.Bytes(), nil
}

func (buf *Memory) UnmarshalBinary(data []byte) error {
	b := bytes.NewBuffer(data)
	_, err := fmt.Fscan(b, &buf.N)
	buf.Buffer = b
	return err
}
<|endoftext|>


<commit_before>package buffer

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"io"
)

type Memory struct {
	N int64
	*bytes.Buffer
}

func New(n int64) *Memory {
	return &Memory{
		N:      n,
		Buffer: bytes.NewBuffer(nil),
	}
}

func (buf *Memory) Cap() int64 {
	return buf.N
}

func (buf *Memory) Len() int64 {
	return int64(buf.Buffer.Len())
}

func (buf *Memory) Write(p []byte) (n int, err error) {
	return LimitWriter(buf.Buffer, Gap(buf)).Write(p)
}

// Must Add io.ErrShortWrite
func (buf *Memory) WriteAt(p []byte, off int64) (n int, err error) {
	if off > buf.Len() {
		return 0, bytes.ErrTooLarge
	} else if len64(p)+off <= buf.Len() {
		d := buf.Bytes()[off:]
		return copy(d, p), nil
	} else {
		d := buf.Bytes()[off:]
		n = copy(d, p)
		m, err := buf.Write(p[n:])
		return n + m, err
	}
}

func (buf *Memory) ReadAt(p []byte, off int64) (n int, err error) {
	return bytes.NewReader(buf.Bytes()).ReadAt(p, off)
}

func (buf *Memory) Read(p []byte) (n int, err error) {
	return io.LimitReader(buf.Buffer, buf.Len()).Read(p)
}

func (buf *Memory) ReadFrom(r io.Reader) (n int64, err error) {
	return buf.Buffer.ReadFrom(io.LimitReader(r, Gap(buf)))
}

func (buf *Memory) FFwd(n int64) int64 {
	data := buf.Bytes()
	if n > len64(data) {
		n = len64(data)
	}
	b := bytes.NewBuffer(data[n:])
	buf.Buffer = b
	return n
}

func init() {
	gob.Register(&Memory{})
}

func (buf *Memory) MarshalBinary() ([]byte, error) {
	var b bytes.Buffer
	fmt.Fprint(&b, buf.N)
	b.Write(buf.Bytes())
	return b.Bytes(), nil
}

func (buf *Memory) UnmarshalBinary(data []byte) error {
	b := bytes.NewBuffer(data)
	_, err := fmt.Fscan(b, &buf.N)
	buf.Buffer = b
	return err
}
<commit_msg>Added support for WriteTo to Memory<commit_after>package buffer

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"io"
)

type Memory struct {
	N int64
	*bytes.Buffer
}

func New(n int64) *Memory {
	return &Memory{
		N:      n,
		Buffer: bytes.NewBuffer(nil),
	}
}

func (buf *Memory) Cap() int64 {
	return buf.N
}

func (buf *Memory) Len() int64 {
	return int64(buf.Buffer.Len())
}

func (buf *Memory) Write(p []byte) (n int, err error) {
	return LimitWriter(buf.Buffer, Gap(buf)).Write(p)
}

func (buf *Memory) WriteTo(w io.Writer) (n int64, err error) {
	return buf.Buffer.WriteTo(w)
}

// Must Add io.ErrShortWrite
func (buf *Memory) WriteAt(p []byte, off int64) (n int, err error) {
	if off > buf.Len() {
		return 0, bytes.ErrTooLarge
	} else if len64(p)+off <= buf.Len() {
		d := buf.Bytes()[off:]
		return copy(d, p), nil
	} else {
		d := buf.Bytes()[off:]
		n = copy(d, p)
		m, err := buf.Write(p[n:])
		return n + m, err
	}
}

func (buf *Memory) ReadAt(p []byte, off int64) (n int, err error) {
	return bytes.NewReader(buf.Bytes()).ReadAt(p, off)
}

func (buf *Memory) Read(p []byte) (n int, err error) {
	return io.LimitReader(buf.Buffer, buf.Len()).Read(p)
}

func (buf *Memory) ReadFrom(r io.Reader) (n int64, err error) {
	return buf.Buffer.ReadFrom(io.LimitReader(r, Gap(buf)))
}

func (buf *Memory) FFwd(n int64) int64 {
	data := buf.Bytes()
	if n > len64(data) {
		n = len64(data)
	}
	b := bytes.NewBuffer(data[n:])
	buf.Buffer = b
	return n
}

func init() {
	gob.Register(&Memory{})
}

func (buf *Memory) MarshalBinary() ([]byte, error) {
	var b bytes.Buffer
	fmt.Fprint(&b, buf.N)
	b.Write(buf.Bytes())
	return b.Bytes(), nil
}

func (buf *Memory) UnmarshalBinary(data []byte) error {
	b := bytes.NewBuffer(data)
	_, err := fmt.Fscan(b, &buf.N)
	buf.Buffer = b
	return err
}
<|endoftext|>


<commit_before>// Package mpd implements parsing and generating of MPEG-DASH Media Presentation Description (MPD) files.
package mpd

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"io"
	"regexp"
	"strconv"
)

// http://mpeg.chiariglione.org/standards/mpeg-dash
// https://www.brendanlong.com/the-structure-of-an-mpeg-dash-mpd.html
// http://standards.iso.org/ittf/PubliclyAvailableStandards/MPEG-DASH_schema_files/DASH-MPD.xsd

var emptyElementRE = regexp.MustCompile(`></[A-Za-z]+>`)

// ConditionalUint (ConditionalUintType) defined in XSD as a union of unsignedInt and boolean.
type ConditionalUint struct {
	u *uint64
	b *bool
}

// MarshalXMLAttr encodes ConditionalUint.
func (c ConditionalUint) MarshalXMLAttr(name xml.Name) (xml.Attr, error) {
	if c.u != nil {
		return xml.Attr{Name: name, Value: strconv.FormatUint(*c.u, 10)}, nil
	}

	if c.b != nil {
		return xml.Attr{Name: name, Value: strconv.FormatBool(*c.b)}, nil
	}

	// both are nil - no attribute, client will threat it like "false"
	return xml.Attr{}, nil
}

// UnmarshalXMLAttr decodes ConditionalUint.
func (c *ConditionalUint) UnmarshalXMLAttr(attr xml.Attr) error {
	u, err := strconv.ParseUint(attr.Value, 10, 64)
	if err == nil {
		c.u = &u
		return nil
	}

	b, err := strconv.ParseBool(attr.Value)
	if err == nil {
		c.b = &b
		return nil
	}

	return fmt.Errorf("ConditionalUint: can't UnmarshalXMLAttr %#v", attr)
}

// check interfaces
var (
	_ xml.MarshalerAttr   = ConditionalUint{}
	_ xml.UnmarshalerAttr = &ConditionalUint{}
)

// MPD represents root XML element.
type MPD struct {
	XMLNS                      *string `xml:"xmlns,attr"`
	Type                       *string `xml:"type,attr"`
	MinimumUpdatePeriod        *string `xml:"minimumUpdatePeriod,attr"`
	AvailabilityStartTime      *string `xml:"availabilityStartTime,attr"`
	MediaPresentationDuration  *string `xml:"mediaPresentationDuration,attr"`
	MinBufferTime              *string `xml:"minBufferTime,attr"`
	SuggestedPresentationDelay *string `xml:"suggestedPresentationDelay,attr"`
	TimeShiftBufferDepth       *string `xml:"timeShiftBufferDepth,attr"`
	Profiles                   string  `xml:"profiles,attr"`
	Period                     *Period `xml:"Period,omitempty"`
}

// Do not try to use encoding.TextMarshaler and encoding.TextUnmarshaler:
// https://github.com/golang/go/issues/6859#issuecomment-118890463

// Encode generates MPD XML.
func (m *MPD) Encode() ([]byte, error) {
	x := new(bytes.Buffer)
	e := xml.NewEncoder(x)
	e.Indent("", "  ")
	err := e.Encode(m)
	if err != nil {
		return nil, err
	}

	// hacks for self-closing tags
	res := new(bytes.Buffer)
	res.WriteString(`<?xml version="1.0" encoding="utf-8"?>`)
	res.WriteByte('\n')
	for {
		s, err := x.ReadString('\n')
		if s != "" {
			s = emptyElementRE.ReplaceAllString(s, `/>`)
			res.WriteString(s)
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}
	}
	res.WriteByte('\n')
	return res.Bytes(), err
}

// Decode parses MPD XML.
func (m *MPD) Decode(b []byte) error {
	return xml.Unmarshal(b, m)
}

// Period represents XSD's PeriodType.
type Period struct {
	Start          *string          `xml:"start,attr"`
	ID             *string          `xml:"id,attr"`
	Duration       *string          `xml:"duration,attr"`
	AdaptationSets []*AdaptationSet `xml:"AdaptationSet,omitempty"`
}

// AdaptationSet represents XSD's AdaptationSetType.
type AdaptationSet struct {
	MimeType                string           `xml:"mimeType,attr"`
	SegmentAlignment        ConditionalUint  `xml:"segmentAlignment,attr"`
	SubsegmentAlignment     ConditionalUint  `xml:"subsegmentAlignment,attr"`
	StartWithSAP            *uint64          `xml:"startWithSAP,attr"`
	SubsegmentStartsWithSAP *uint64          `xml:"subsegmentStartsWithSAP,attr"`
	BitstreamSwitching      *bool            `xml:"bitstreamSwitching,attr"`
	Lang                    *string          `xml:"lang,attr"`
	ContentProtections      []Descriptor     `xml:"ContentProtection,omitempty"`
	Representations         []Representation `xml:"Representation,omitempty"`
}

// Representation represents XSD's RepresentationType.
type Representation struct {
	ID                 *string          `xml:"id,attr"`
	Width              *uint64          `xml:"width,attr"`
	Height             *uint64          `xml:"height,attr"`
	FrameRate          *string          `xml:"frameRate,attr"`
	Bandwidth          *uint64          `xml:"bandwidth,attr"`
	AudioSamplingRate  *string          `xml:"audioSamplingRate,attr"`
	Codecs             *string          `xml:"codecs,attr"`
	ContentProtections []Descriptor     `xml:"ContentProtection,omitempty"`
	SegmentTemplate    *SegmentTemplate `xml:"SegmentTemplate,omitempty"`
}

// Descriptor represents XSD's DescriptorType.
type Descriptor struct {
	SchemeIDURI *string `xml:"schemeIdUri,attr"`
	Value       *string `xml:"value,attr"`
}

// SegmentTemplate represents XSD's SegmentTemplateType.
type SegmentTemplate struct {
	Timescale              *uint64            `xml:"timescale,attr"`
	Media                  *string            `xml:"media,attr"`
	Initialization         *string            `xml:"initialization,attr"`
	StartNumber            *uint64            `xml:"startNumber,attr"`
	PresentationTimeOffset *uint64            `xml:"presentationTimeOffset,attr"`
	SegmentTimelineS       []SegmentTimelineS `xml:"SegmentTimeline>S,omitempty"`
}

// SegmentTimelineS represents XSD's SegmentTimelineType's inner S elements.
type SegmentTimelineS struct {
	T *uint64 `xml:"t,attr"`
	D uint64  `xml:"d,attr"`
	R *int64  `xml:"r,attr"`
}
<commit_msg>add publishTime attribute<commit_after>// Package mpd implements parsing and generating of MPEG-DASH Media Presentation Description (MPD) files.
package mpd

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"io"
	"regexp"
	"strconv"
)

// http://mpeg.chiariglione.org/standards/mpeg-dash
// https://www.brendanlong.com/the-structure-of-an-mpeg-dash-mpd.html
// http://standards.iso.org/ittf/PubliclyAvailableStandards/MPEG-DASH_schema_files/DASH-MPD.xsd

var emptyElementRE = regexp.MustCompile(`></[A-Za-z]+>`)

// ConditionalUint (ConditionalUintType) defined in XSD as a union of unsignedInt and boolean.
type ConditionalUint struct {
	u *uint64
	b *bool
}

// MarshalXMLAttr encodes ConditionalUint.
func (c ConditionalUint) MarshalXMLAttr(name xml.Name) (xml.Attr, error) {
	if c.u != nil {
		return xml.Attr{Name: name, Value: strconv.FormatUint(*c.u, 10)}, nil
	}

	if c.b != nil {
		return xml.Attr{Name: name, Value: strconv.FormatBool(*c.b)}, nil
	}

	// both are nil - no attribute, client will threat it like "false"
	return xml.Attr{}, nil
}

// UnmarshalXMLAttr decodes ConditionalUint.
func (c *ConditionalUint) UnmarshalXMLAttr(attr xml.Attr) error {
	u, err := strconv.ParseUint(attr.Value, 10, 64)
	if err == nil {
		c.u = &u
		return nil
	}

	b, err := strconv.ParseBool(attr.Value)
	if err == nil {
		c.b = &b
		return nil
	}

	return fmt.Errorf("ConditionalUint: can't UnmarshalXMLAttr %#v", attr)
}

// check interfaces
var (
	_ xml.MarshalerAttr   = ConditionalUint{}
	_ xml.UnmarshalerAttr = &ConditionalUint{}
)

// MPD represents root XML element.
type MPD struct {
	XMLNS                      *string `xml:"xmlns,attr"`
	Type                       *string `xml:"type,attr"`
	MinimumUpdatePeriod        *string `xml:"minimumUpdatePeriod,attr"`
	AvailabilityStartTime      *string `xml:"availabilityStartTime,attr"`
	MediaPresentationDuration  *string `xml:"mediaPresentationDuration,attr"`
	MinBufferTime              *string `xml:"minBufferTime,attr"`
	SuggestedPresentationDelay *string `xml:"suggestedPresentationDelay,attr"`
	TimeShiftBufferDepth       *string `xml:"timeShiftBufferDepth,attr"`
	PublishTime                *string `xml:"publishTime,attr"`
	Profiles                   string  `xml:"profiles,attr"`
	Period                     *Period `xml:"Period,omitempty"`
}

// Do not try to use encoding.TextMarshaler and encoding.TextUnmarshaler:
// https://github.com/golang/go/issues/6859#issuecomment-118890463

// Encode generates MPD XML.
func (m *MPD) Encode() ([]byte, error) {
	x := new(bytes.Buffer)
	e := xml.NewEncoder(x)
	e.Indent("", "  ")
	err := e.Encode(m)
	if err != nil {
		return nil, err
	}

	// hacks for self-closing tags
	res := new(bytes.Buffer)
	res.WriteString(`<?xml version="1.0" encoding="utf-8"?>`)
	res.WriteByte('\n')
	for {
		s, err := x.ReadString('\n')
		if s != "" {
			s = emptyElementRE.ReplaceAllString(s, `/>`)
			res.WriteString(s)
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}
	}
	res.WriteByte('\n')
	return res.Bytes(), err
}

// Decode parses MPD XML.
func (m *MPD) Decode(b []byte) error {
	return xml.Unmarshal(b, m)
}

// Period represents XSD's PeriodType.
type Period struct {
	Start          *string          `xml:"start,attr"`
	ID             *string          `xml:"id,attr"`
	Duration       *string          `xml:"duration,attr"`
	AdaptationSets []*AdaptationSet `xml:"AdaptationSet,omitempty"`
}

// AdaptationSet represents XSD's AdaptationSetType.
type AdaptationSet struct {
	MimeType                string           `xml:"mimeType,attr"`
	SegmentAlignment        ConditionalUint  `xml:"segmentAlignment,attr"`
	SubsegmentAlignment     ConditionalUint  `xml:"subsegmentAlignment,attr"`
	StartWithSAP            *uint64          `xml:"startWithSAP,attr"`
	SubsegmentStartsWithSAP *uint64          `xml:"subsegmentStartsWithSAP,attr"`
	BitstreamSwitching      *bool            `xml:"bitstreamSwitching,attr"`
	Lang                    *string          `xml:"lang,attr"`
	ContentProtections      []Descriptor     `xml:"ContentProtection,omitempty"`
	Representations         []Representation `xml:"Representation,omitempty"`
}

// Representation represents XSD's RepresentationType.
type Representation struct {
	ID                 *string          `xml:"id,attr"`
	Width              *uint64          `xml:"width,attr"`
	Height             *uint64          `xml:"height,attr"`
	FrameRate          *string          `xml:"frameRate,attr"`
	Bandwidth          *uint64          `xml:"bandwidth,attr"`
	AudioSamplingRate  *string          `xml:"audioSamplingRate,attr"`
	Codecs             *string          `xml:"codecs,attr"`
	ContentProtections []Descriptor     `xml:"ContentProtection,omitempty"`
	SegmentTemplate    *SegmentTemplate `xml:"SegmentTemplate,omitempty"`
}

// Descriptor represents XSD's DescriptorType.
type Descriptor struct {
	SchemeIDURI *string `xml:"schemeIdUri,attr"`
	Value       *string `xml:"value,attr"`
}

// SegmentTemplate represents XSD's SegmentTemplateType.
type SegmentTemplate struct {
	Timescale              *uint64            `xml:"timescale,attr"`
	Media                  *string            `xml:"media,attr"`
	Initialization         *string            `xml:"initialization,attr"`
	StartNumber            *uint64            `xml:"startNumber,attr"`
	PresentationTimeOffset *uint64            `xml:"presentationTimeOffset,attr"`
	SegmentTimelineS       []SegmentTimelineS `xml:"SegmentTimeline>S,omitempty"`
}

// SegmentTimelineS represents XSD's SegmentTimelineType's inner S elements.
type SegmentTimelineS struct {
	T *uint64 `xml:"t,attr"`
	D uint64  `xml:"d,attr"`
	R *int64  `xml:"r,attr"`
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"time"
)

import "github.com/JamesDunne/go-util/base"

const buffer_size = 4096
const debug = false

type conn struct {
	c      net.Conn
	logger *log.Logger
	buffer []byte

	packet0 [][]byte
}

func newConn(c net.Conn) *conn {
	return &conn{
		c:       c,
		buffer:  make([]byte, buffer_size),
		logger:  log.New(os.Stderr, fmt.Sprintf("%s: ", c.RemoteAddr()), 0),
		packet0: make([][]byte, 0, 5),
	}
}

// 200ms timeout for SSH detection:
const timeoutDuration = time.Millisecond * time.Duration(500)

func (c *conn) xmit(r io.Reader, w io.Writer, stop chan bool) {
	buffer := make([]byte, buffer_size)
	for {
		rn, err := r.Read(buffer)
		if err != nil {
			if err == io.EOF {
				return
			}
			if verbose {
				c.logger.Println(err)
			}
			break
		}

		_, err = w.Write(buffer[:rn])
		if err != nil {
			if err == io.EOF {
				return
			}
			if verbose {
				c.logger.Println(err)
			}
			break
		}
	}

	stop <- true
}

// Handles a single connection and sniffs the protocol:
func (c *conn) serve() {
	client := c.c

	if verbose {
		defer c.logger.Println("closed")
	}
	defer client.Close()

	if verbose {
		c.logger.Println("accepted")
	}

	var target_addr *base.Dialable
	sniffed := false
	for !sniffed {
		// Set a timeout on sniffing because some SSH clients (PuTTY) will wait eternally for incoming data before sending
		// their first packets:
		client.SetReadDeadline(time.Now().Add(timeoutDuration))

		// Read some data:
		n, err := client.Read(c.buffer)
		if err == io.EOF {
			return
		}
		if _, ok := err.(net.Error); ok {
			// Timed out; assume SSH:
			if verbose {
				c.logger.Println("timed out; assuming SSH")
			}
			sniffed = true
			target_addr = ssh_addr
			break
		}
		if err != nil {
			return
		}

		p := c.buffer[0:n]
		if debug && verbose {
			base.HexDumpToLogger(p, c.logger)
		}

		// Keep a rolling buffer of the packets we're sniffing so we can send them later:
		c.packet0 = append(c.packet0, p)

		// Check if TLS protocol:
		if n < 3 {
			continue
		}
		// TLS packet starts with a record "Hello" (0x16), followed by version (0x03 0x00-0x03) (RFC6101 A.1)
		// Reject SSLv2 and lower versions (RFC6176):
		if p[0] == 0x16 && p[1] == 0x03 && (p[2] >= 0x00 && p[2] <= 0x03) {
			sniffed = true
			target_addr = https_addr
			if verbose {
				c.logger.Println("detected HTTPS")
			}
			break
		}

		// Check if SSH protocol:
		if n < 4 {
			continue
		}
		if p[0] == 'S' && p[1] == 'S' && p[2] == 'H' && p[3] == '-' {
			sniffed = true
			target_addr = ssh_addr
			if verbose {
				c.logger.Println("detected SSH")
			}
			break
		}
	}

	// Clear the deadline:
	client.SetReadDeadline(time.Time{})

	// Now just copy data from in to out:
	server, err := net.Dial(target_addr.Network, target_addr.Address)
	if err != nil {
		if verbose {
			c.logger.Printf("%s\n", err)
		}
		return
	}

	// Transmit first packet(s) that we sniffed:
	for _, p := range c.packet0 {
		_, err = server.Write(p)
		if err == io.EOF {
			return
		}
	}

	// Start proxying traffic both ways:
	complete := make(chan bool, 1)

	// From client to server:
	go c.xmit(client, server, complete)
	go c.xmit(server, client, complete)

	// Wait until either end closes:
	<-complete
}

func serveMux(l net.Listener) (err error) {
	defer l.Close()
	var tempDelay time.Duration // how long to sleep on accept failure
	for {
		rw, e := l.Accept()
		if e != nil {
			if ne, ok := e.(net.Error); ok && ne.Temporary() {
				if tempDelay == 0 {
					tempDelay = 5 * time.Millisecond
				} else {
					tempDelay *= 2
				}
				if max := 1 * time.Second; tempDelay > max {
					tempDelay = max
				}
				log.Printf("Accept error: %v; retrying in %v\n", e, tempDelay)
				time.Sleep(tempDelay)
				continue
			}
			return e
		}
		tempDelay = 0

		c := newConn(rw)
		// Launch a goroutine to handle traffic:
		go c.serve()
	}

	return nil
}
<commit_msg>Forgot to close server connection.<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"time"
)

import "github.com/JamesDunne/go-util/base"

const buffer_size = 4096
const debug = false

type conn struct {
	c      net.Conn
	logger *log.Logger
	buffer []byte

	packet0 [][]byte
}

func newConn(c net.Conn) *conn {
	return &conn{
		c:       c,
		buffer:  make([]byte, buffer_size),
		logger:  log.New(os.Stderr, fmt.Sprintf("%s: ", c.RemoteAddr()), 0),
		packet0: make([][]byte, 0, 5),
	}
}

// 200ms timeout for SSH detection:
const timeoutDuration = time.Millisecond * time.Duration(500)

func (c *conn) xmit(r io.Reader, w io.Writer, stop chan bool) {
	buffer := make([]byte, buffer_size)
	for {
		rn, err := r.Read(buffer)
		if err != nil {
			if err == io.EOF {
				return
			}
			if verbose {
				c.logger.Println(err)
			}
			break
		}

		_, err = w.Write(buffer[:rn])
		if err != nil {
			if err == io.EOF {
				return
			}
			if verbose {
				c.logger.Println(err)
			}
			break
		}
	}

	stop <- true
}

// Handles a single connection and sniffs the protocol:
func (c *conn) serve() {
	client := c.c

	if verbose {
		defer c.logger.Println("closed")
	}
	defer client.Close()

	if verbose {
		c.logger.Println("accepted")
	}

	var target_addr *base.Dialable
	sniffed := false
	for !sniffed {
		// Set a timeout on sniffing because some SSH clients (PuTTY) will wait eternally for incoming data before sending
		// their first packets:
		client.SetReadDeadline(time.Now().Add(timeoutDuration))

		// Read some data:
		n, err := client.Read(c.buffer)
		if err == io.EOF {
			return
		}
		if _, ok := err.(net.Error); ok {
			// Timed out; assume SSH:
			if verbose {
				c.logger.Println("timed out; assuming SSH")
			}
			sniffed = true
			target_addr = ssh_addr
			break
		}
		if err != nil {
			return
		}

		p := c.buffer[0:n]
		if debug && verbose {
			base.HexDumpToLogger(p, c.logger)
		}

		// Keep a rolling buffer of the packets we're sniffing so we can send them later:
		c.packet0 = append(c.packet0, p)

		// Check if TLS protocol:
		if n < 3 {
			continue
		}
		// TLS packet starts with a record "Hello" (0x16), followed by version (0x03 0x00-0x03) (RFC6101 A.1)
		// Reject SSLv2 and lower versions (RFC6176):
		if p[0] == 0x16 && p[1] == 0x03 && (p[2] >= 0x00 && p[2] <= 0x03) {
			sniffed = true
			target_addr = https_addr
			if verbose {
				c.logger.Println("detected HTTPS")
			}
			break
		}

		// Check if SSH protocol:
		if n < 4 {
			continue
		}
		if p[0] == 'S' && p[1] == 'S' && p[2] == 'H' && p[3] == '-' {
			sniffed = true
			target_addr = ssh_addr
			if verbose {
				c.logger.Println("detected SSH")
			}
			break
		}
	}

	// Clear the deadline:
	client.SetReadDeadline(time.Time{})

	// Now just copy data from in to out:
	server, err := net.Dial(target_addr.Network, target_addr.Address)
	if err != nil {
		if verbose {
			c.logger.Printf("%s\n", err)
		}
		return
	}
	defer server.Close()

	// Transmit first packet(s) that we sniffed:
	for _, p := range c.packet0 {
		_, err = server.Write(p)
		if err == io.EOF {
			return
		}
	}

	// Start proxying traffic both ways:
	complete := make(chan bool, 1)

	// From client to server:
	go c.xmit(client, server, complete)
	go c.xmit(server, client, complete)

	// Wait until either end closes:
	<-complete
}

func serveMux(l net.Listener) (err error) {
	defer l.Close()
	var tempDelay time.Duration // how long to sleep on accept failure
	for {
		rw, e := l.Accept()
		if e != nil {
			if ne, ok := e.(net.Error); ok && ne.Temporary() {
				if tempDelay == 0 {
					tempDelay = 5 * time.Millisecond
				} else {
					tempDelay *= 2
				}
				if max := 1 * time.Second; tempDelay > max {
					tempDelay = max
				}
				log.Printf("Accept error: %v; retrying in %v\n", e, tempDelay)
				time.Sleep(tempDelay)
				continue
			}
			return e
		}
		tempDelay = 0

		c := newConn(rw)
		// Launch a goroutine to handle traffic:
		go c.serve()
	}

	return nil
}
<|endoftext|>


<commit_before>package ndp

import (
    "encoding/binary"
    "errors"
)

func ParseRouterAdvertisement(b []byte) (*RouterAdvertisement, error) {
    if len(b) < 4 {
        return nil, errors.New("message too short")
    }

    // first 4 bytes is ICMPv6 header
    p := &RouterAdvertisement{
        HopLimit:       uint8(b[4]),
        ManagedAddress: false,
        OtherStateful:  false,
        HomeAgent:      false,
        RouterLifeTime: binary.BigEndian.Uint16(b[6:8]),
        ReachableTime:  binary.BigEndian.Uint32(b[8:12]),
        RetransTimer:   binary.BigEndian.Uint32(b[12:16]),
    }

    // parse flags
    if b[5] & 0x80 > 0 {
        p.ManagedAddress = true
    }
    if b[5] & 0x40 > 0 {
        p.OtherStateful = true
    }
    if b[5] & 0x20 > 0 {
        p.HomeAgent = true
    }

    return p, nil
}

func ParseRouterSolicitation(b []byte) (*RouterSolicitation, error) {
    if len(b) < 4 {
        return nil, errors.New("message too short")
    }

    // first 4 bytes is ICMPv6 header
    p := &RouterSolicitation{

    }

    return p, nil
}

func ParseNeighborAdvertisement(b []byte) (*NeighborAdvertisement, error) {
    if len(b) < 4 {
        return nil, errors.New("message too short")
    }

    // first 4 bytes is ICMPv6 header
    p := &NeighborAdvertisement{
       TargetAddress: b[8:24],
    }

    return p, nil
}

func ParseNeighborSolicitation(b []byte) (*NeighborSolicitation, error) {
    if len(b) < 4 {
        return nil, errors.New("message too short")
    }

    // first 4 bytes is ICMPv6 header
    p := &NeighborSolicitation{
       TargetAddress: b[8:24],
    }

    return p, nil
}
<commit_msg>parse NA flags<commit_after>package ndp

import (
    "encoding/binary"
    "errors"
)

func ParseRouterAdvertisement(b []byte) (*RouterAdvertisement, error) {
    if len(b) < 4 {
        return nil, errors.New("message too short")
    }

    // first 4 bytes is ICMPv6 header
    p := &RouterAdvertisement{
        HopLimit:       uint8(b[4]),
        ManagedAddress: false,
        OtherStateful:  false,
        HomeAgent:      false,
        RouterLifeTime: binary.BigEndian.Uint16(b[6:8]),
        ReachableTime:  binary.BigEndian.Uint32(b[8:12]),
        RetransTimer:   binary.BigEndian.Uint32(b[12:16]),
    }

    // parse flags
    if b[5] & 0x80 > 0 {
        p.ManagedAddress = true
    }
    if b[5] & 0x40 > 0 {
        p.OtherStateful = true
    }
    if b[5] & 0x20 > 0 {
        p.HomeAgent = true
    }

    return p, nil
}

func ParseRouterSolicitation(b []byte) (*RouterSolicitation, error) {
    if len(b) < 4 {
        return nil, errors.New("message too short")
    }

    // first 4 bytes is ICMPv6 header
    p := &RouterSolicitation{

    }

    return p, nil
}

func ParseNeighborAdvertisement(b []byte) (*NeighborAdvertisement, error) {
    if len(b) < 4 {
        return nil, errors.New("message too short")
    }

    // first 4 bytes is ICMPv6 header
    p := &NeighborAdvertisement{
       TargetAddress: b[8:24],
    }

    // parse flags
    if b[4] & 0x80 > 0 {
        p.Router = true
    }
    if b[4] & 0x40 > 0 {
        p.Solicited = true
    }
    if b[4] & 0x20 > 0 {
        p.Override = true
    }

    return p, nil
}

func ParseNeighborSolicitation(b []byte) (*NeighborSolicitation, error) {
    if len(b) < 4 {
        return nil, errors.New("message too short")
    }

    // first 4 bytes is ICMPv6 header
    p := &NeighborSolicitation{
       TargetAddress: b[8:24],
    }

    return p, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Google 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.

package fuse

import (
	"syscall"
	"unsafe"

	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/fuse/internal/buffer"
	"github.com/jacobsa/fuse/internal/fusekernel"
)

// Given an op, return the response that should be sent to the kernel. If the
// op requires no response, return a nil response. If the op is unknown, return
// an error that should be sent to the kernel.
func kernelResponse(
	untyped fuseops.Op,
	protocol fusekernel.Protocol) (b buffer.OutMessage, err error) {
	switch o := untyped.(type) {
	case *fuseops.LookUpInodeOp:
		size := fusekernel.EntryOutSize(protocol)
		b = buffer.NewOutMessage(size)
		out := (*fusekernel.EntryOut)(b.Grow(size))
		convertChildInodeEntry(&o.Entry, out)

	case *fuseops.GetInodeAttributesOp:
		size := fusekernel.AttrOutSize(protocol)
		b = buffer.NewOutMessage(size)
		out := (*fusekernel.AttrOut)(b.Grow(size))
		out.AttrValid, out.AttrValidNsec = convertExpirationTime(o.AttributesExpiration)
		convertAttributes(o.Inode, &o.Attributes, &out.Attr)

	case *fuseops.SetInodeAttributesOp:
		size := fusekernel.AttrOutSize(protocol)
		b = buffer.NewOutMessage(size)
		out := (*fusekernel.AttrOut)(b.Grow(size))
		out.AttrValid, out.AttrValidNsec = convertExpirationTime(o.AttributesExpiration)
		convertAttributes(o.Inode, &o.Attributes, &out.Attr)

	case *fuseops.ForgetInodeOp:
		// No response.

	case *fuseops.MkDirOp:
		size := fusekernel.EntryOutSize(protocol)
		b = buffer.NewOutMessage(size)
		out := (*fusekernel.EntryOut)(b.Grow(size))
		convertChildInodeEntry(&o.Entry, out)

	case *fuseops.CreateFileOp:
		eSize := fusekernel.EntryOutSize(protocol)
		b = buffer.NewOutMessage(eSize + unsafe.Sizeof(fusekernel.OpenOut{}))

		e := (*fusekernel.EntryOut)(b.Grow(eSize))
		convertChildInodeEntry(&o.Entry, e)

		oo := (*fusekernel.OpenOut)(b.Grow(unsafe.Sizeof(fusekernel.OpenOut{})))
		oo.Fh = uint64(o.Handle)

	case *fuseops.CreateSymlinkOp:
		size := fusekernel.EntryOutSize(protocol)
		b = buffer.NewOutMessage(size)
		out := (*fusekernel.EntryOut)(b.Grow(size))
		convertChildInodeEntry(&o.Entry, out)

	case *fuseops.RenameOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.RmDirOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.UnlinkOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.OpenDirOp:
		b = buffer.NewOutMessage(unsafe.Sizeof(fusekernel.OpenOut{}))
		out := (*fusekernel.OpenOut)(b.Grow(unsafe.Sizeof(fusekernel.OpenOut{})))
		out.Fh = uint64(o.Handle)

	case *fuseops.ReadDirOp:
		b = buffer.NewOutMessage(uintptr(len(o.Data)))
		b.Append(o.Data)

	case *fuseops.ReleaseDirHandleOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.OpenFileOp:
		b = buffer.NewOutMessage(unsafe.Sizeof(fusekernel.OpenOut{}))
		out := (*fusekernel.OpenOut)(b.Grow(unsafe.Sizeof(fusekernel.OpenOut{})))
		out.Fh = uint64(o.Handle)

	case *fuseops.ReadFileOp:
		b = buffer.NewOutMessage(uintptr(len(o.Data)))
		b.Append(o.Data)

	case *fuseops.WriteFileOp:
		b = buffer.NewOutMessage(unsafe.Sizeof(fusekernel.WriteOut{}))
		out := (*fusekernel.WriteOut)(b.Grow(unsafe.Sizeof(fusekernel.WriteOut{})))
		out.Size = uint32(len(o.Data))

	case *fuseops.SyncFileOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.FlushFileOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.ReleaseFileHandleOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.ReadSymlinkOp:
		b = buffer.NewOutMessage(uintptr(len(o.Target)))
		b.AppendString(o.Target)

	case *internalStatFSOp:
		b = buffer.NewOutMessage(unsafe.Sizeof(fusekernel.StatfsOut{}))
		b.Grow(unsafe.Sizeof(fusekernel.StatfsOut{}))

	case *internalInterruptOp:
		// No response.

	case *internalInitOp:
		b = buffer.NewOutMessage(unsafe.Sizeof(fusekernel.InitOut{}))
		out := (*fusekernel.InitOut)(b.Grow(unsafe.Sizeof(fusekernel.InitOut{})))

		out.Major = o.Library.Major
		out.Minor = o.Library.Minor
		out.MaxReadahead = o.MaxReadahead
		out.Flags = uint32(o.Flags)
		out.MaxWrite = o.MaxWrite

	default:
		err = syscall.ENOSYS
	}

	return
}
<commit_msg>Restored the internal ops.<commit_after>// Copyright 2015 Google 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.

package fuse

import (
	"syscall"
	"unsafe"

	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/fuse/internal/buffer"
	"github.com/jacobsa/fuse/internal/fusekernel"
)

// Given an op, return the response that should be sent to the kernel. If the
// op requires no response, return a nil response. If the op is unknown, return
// an error that should be sent to the kernel.
func kernelResponse(
	untyped fuseops.Op,
	protocol fusekernel.Protocol) (b buffer.OutMessage, err error) {
	switch o := untyped.(type) {
	case *fuseops.LookUpInodeOp:
		size := fusekernel.EntryOutSize(protocol)
		b = buffer.NewOutMessage(size)
		out := (*fusekernel.EntryOut)(b.Grow(size))
		convertChildInodeEntry(&o.Entry, out)

	case *fuseops.GetInodeAttributesOp:
		size := fusekernel.AttrOutSize(protocol)
		b = buffer.NewOutMessage(size)
		out := (*fusekernel.AttrOut)(b.Grow(size))
		out.AttrValid, out.AttrValidNsec = convertExpirationTime(o.AttributesExpiration)
		convertAttributes(o.Inode, &o.Attributes, &out.Attr)

	case *fuseops.SetInodeAttributesOp:
		size := fusekernel.AttrOutSize(protocol)
		b = buffer.NewOutMessage(size)
		out := (*fusekernel.AttrOut)(b.Grow(size))
		out.AttrValid, out.AttrValidNsec = convertExpirationTime(o.AttributesExpiration)
		convertAttributes(o.Inode, &o.Attributes, &out.Attr)

	case *fuseops.ForgetInodeOp:
		// No response.

	case *fuseops.MkDirOp:
		size := fusekernel.EntryOutSize(protocol)
		b = buffer.NewOutMessage(size)
		out := (*fusekernel.EntryOut)(b.Grow(size))
		convertChildInodeEntry(&o.Entry, out)

	case *fuseops.CreateFileOp:
		eSize := fusekernel.EntryOutSize(protocol)
		b = buffer.NewOutMessage(eSize + unsafe.Sizeof(fusekernel.OpenOut{}))

		e := (*fusekernel.EntryOut)(b.Grow(eSize))
		convertChildInodeEntry(&o.Entry, e)

		oo := (*fusekernel.OpenOut)(b.Grow(unsafe.Sizeof(fusekernel.OpenOut{})))
		oo.Fh = uint64(o.Handle)

	case *fuseops.CreateSymlinkOp:
		size := fusekernel.EntryOutSize(protocol)
		b = buffer.NewOutMessage(size)
		out := (*fusekernel.EntryOut)(b.Grow(size))
		convertChildInodeEntry(&o.Entry, out)

	case *fuseops.RenameOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.RmDirOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.UnlinkOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.OpenDirOp:
		b = buffer.NewOutMessage(unsafe.Sizeof(fusekernel.OpenOut{}))
		out := (*fusekernel.OpenOut)(b.Grow(unsafe.Sizeof(fusekernel.OpenOut{})))
		out.Fh = uint64(o.Handle)

	case *fuseops.ReadDirOp:
		b = buffer.NewOutMessage(uintptr(len(o.Data)))
		b.Append(o.Data)

	case *fuseops.ReleaseDirHandleOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.OpenFileOp:
		b = buffer.NewOutMessage(unsafe.Sizeof(fusekernel.OpenOut{}))
		out := (*fusekernel.OpenOut)(b.Grow(unsafe.Sizeof(fusekernel.OpenOut{})))
		out.Fh = uint64(o.Handle)

	case *fuseops.ReadFileOp:
		b = buffer.NewOutMessage(uintptr(len(o.Data)))
		b.Append(o.Data)

	case *fuseops.WriteFileOp:
		b = buffer.NewOutMessage(unsafe.Sizeof(fusekernel.WriteOut{}))
		out := (*fusekernel.WriteOut)(b.Grow(unsafe.Sizeof(fusekernel.WriteOut{})))
		out.Size = uint32(len(o.Data))

	case *fuseops.SyncFileOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.FlushFileOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.ReleaseFileHandleOp:
		b = buffer.NewOutMessage(0)

	case *fuseops.ReadSymlinkOp:
		b = buffer.NewOutMessage(uintptr(len(o.Target)))
		b.AppendString(o.Target)

	case *internalStatFSOp:
		b = buffer.NewOutMessage(unsafe.Sizeof(fusekernel.StatfsOut{}))
		b.Grow(unsafe.Sizeof(fusekernel.StatfsOut{}))

	case *internalInterruptOp:
		// No response.

	case *internalInitOp:
		b = buffer.NewOutMessage(unsafe.Sizeof(fusekernel.InitOut{}))
		out := (*fusekernel.InitOut)(b.Grow(unsafe.Sizeof(fusekernel.InitOut{})))

		out.Major = o.Library.Major
		out.Minor = o.Library.Minor
		out.MaxReadahead = o.MaxReadahead
		out.Flags = uint32(o.Flags)
		out.MaxWrite = o.MaxWrite

	default:
		err = syscall.ENOSYS
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Internal
////////////////////////////////////////////////////////////////////////

// Common implementation for our "internal" ops that don't need to be pretty.
type internalOp struct {
}

func (o *internalOp) ShortDesc() string   { return "<internalOp>" }
func (o *internalOp) DebugString() string { return "<internalOp>" }

type internalStatFSOp struct {
	internalOp
}

func (o *internalStatFSOp) kernelResponse() (b buffer.OutMessage) {
	b = buffer.NewOutMessage(unsafe.Sizeof(fusekernel.StatfsOut{}))
	b.Grow(unsafe.Sizeof(fusekernel.StatfsOut{}))

	return
}

type internalInterruptOp struct {
	internalOp
	FuseID uint64
}

func (o *internalInterruptOp) kernelResponse() (b buffer.OutMessage) {
	panic("Shouldn't get here.")
}

type internalInitOp struct {
	internalOp

	// In
	Kernel fusekernel.Protocol

	// Out
	Library      fusekernel.Protocol
	MaxReadahead uint32
	Flags        fusekernel.InitFlags
	MaxWrite     uint32
}

func (o *internalInitOp) kernelResponse() (b buffer.OutMessage) {
	b = buffer.NewOutMessage(unsafe.Sizeof(fusekernel.InitOut{}))
	out := (*fusekernel.InitOut)(b.Grow(unsafe.Sizeof(fusekernel.InitOut{})))

	out.Major = o.Library.Major
	out.Minor = o.Library.Minor
	out.MaxReadahead = o.MaxReadahead
	out.Flags = uint32(o.Flags)
	out.MaxWrite = o.MaxWrite

	return
}
<|endoftext|>


<commit_before>package pbs

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/golang/protobuf/proto"
	"io"
	"reflect"
	"strconv"
	"strings"
)

const (
	Varint      = 0
	Int64       = 1
	LengthDelim = 2
	StartGroup  = 3
	EndGroup    = 4
	Bit32       = 5
)

type StreamMessage interface {
	proto.Message
	io.Closer
	Errors() chan error
	Closed() <-chan struct{}
}

func splitTypeAndField(b byte) (Type, Field byte) {
	return (b & 0x7), b >> 3
}

func readLengthDelim(r *bufio.Reader) ([]byte, error) {
	l, err := readVarint(r)
	if err != nil {
		return nil, err
	}

	buf := make([]byte, l)
	_, err = io.ReadFull(r, buf)
	if err != nil {
		return nil, err
	}
	return buf, nil
}

func readVarint(r *bufio.Reader) (int, error) {
	var sum int
	for i := uint(0); i < 4; i++ {
		b, err := r.ReadByte()
		if err != nil {
			return 0, err
		}

		cont := b & 128
		val := b & 127

		sum += int(val << (7 * i))
		if cont == 0 {
			break
		}
	}
	return sum, nil
}

func (db *decBuffer) decodeField(field byte, data []byte) error {
	finfo := db.props.FieldMapping[field]

	val := reflect.ValueOf(db.val).Elem()

	f := val.Field(finfo.GoField)
	switch {
	case f.Type().Kind() == reflect.Chan:
		// This is a 'repeated' field
		// we just encode this value and send it along

		e := f.Type().Elem()
		if e.Kind() == reflect.Ptr {
			e = e.Elem()
		}

		newVal := reflect.New(e)
		switch pv := newVal.Interface().(type) {
		case proto.Message:
			err := proto.Unmarshal(data, pv)
			if err != nil {
				return err
			}
			f.Send(newVal)
		case *string:
			*pv = string(data)
			f.Send(newVal.Elem())
		case *[]byte:
			*pv = data
			f.Send(newVal.Elem())
		default:
			fmt.Println(reflect.TypeOf(pv))
			fmt.Println(data)
			return fmt.Errorf("Unrecognized type in protobuf field decode")
		}
	case f.Type().Elem().Kind() == reflect.String:
		f.Set(reflect.New(f.Type().Elem()))
		f.Elem().SetString(string(data))
	default:
		fmt.Println("UNKNOWN")
		fmt.Println(f)
		return errors.New("TODO: fix this unsupported type")
	}
	return nil
}

type decBuffer struct {
	props *Props
	val   StreamMessage
}

func StreamDecode(r io.Reader, sm StreamMessage) error {
	props, err := GetProperties(sm)
	if err != nil {
		return err
	}

	go func() {
		defer sm.Close()
		read := bufio.NewReader(r)

		db := decBuffer{
			props: props,
			val:   sm,
		}
		for {
			b, err := read.ReadByte()
			if err != nil {
				if err == io.EOF {
					return
				}
				sm.Errors() <- err
				return
			}

			typ, f := splitTypeAndField(b)
			switch typ {
			case Varint:
				i, err := readVarint(read)
				if err != nil {
					sm.Errors() <- err
					return
				}

				field := reflect.ValueOf(sm).Elem().Field(props.FieldMapping[f].GoField)

				fmt.Println("proto field: ", f)
				fmt.Println("FIELD: ", field)
				elemType := field.Type()
				if elemType.Kind() == reflect.Chan {
					elemType = elemType.Elem()
				}
				fmt.Println("elemtype: ", elemType)

				nval := reflect.New(elemType.Elem())
				nval.Elem().SetInt(int64(i))

				if props.FieldMapping[f].Repeated {
					field.Send(nval.Elem())
				} else {
					field.Set(nval)
				}
			case LengthDelim:
				val, err := readLengthDelim(read)
				if err != nil {
					if err == io.EOF {
						return
					}
					sm.Errors() <- err
					return
				}
				err = db.decodeField(f, val)
				if err != nil {
					sm.Errors() <- err
					return
				}
			default:
				fmt.Println("unrecognized type: ", typ)
				return
			}
		}
	}()
	return nil
}

func combineTypeAndField(typ, field byte) byte {
	return (field << 3) | (typ & 0x7)
}

func writeLengthDelimited(w io.Writer, data []byte) error {
	length := proto.EncodeVarint(uint64(len(data)))
	n, err := w.Write(length)
	if err != nil {
		return err
	}
	if n != len(length) {
		return errors.New("failed to write length")
	}

	n, err = w.Write([]byte(data))
	if err != nil {
		return err
	}
	if n != len(data) {
		return errors.New("failed to write data")
	}
	return nil
}

func writeTag(w io.Writer, tag byte) error {
	n, err := w.Write([]byte{tag})
	if err != nil {
		return err
	}
	if n != 1 {
		return errors.New("failed to write tag")
	}
	return nil
}

func writeVarint(w io.Writer, v uint64) error {
	data := proto.EncodeVarint(v)
	n, err := w.Write(data)
	if err != nil {
		return err
	}
	if n != len(data) {
		return errors.New("failed to write enough bytes")
	}
	return nil
}

func StreamEncode(w io.Writer, sm StreamMessage) error {
	props, err := GetProperties(sm)
	if err != nil {
		return err
	}

	val := reflect.ValueOf(sm).Elem()

	for protoField, fprop := range props.FieldMapping {
		field := val.Field(fprop.GoField)
		val := field
		if val.Kind() == reflect.Ptr {
			val = val.Elem()
		}
		if !fprop.Repeated {
			switch val.Kind() {
			case reflect.String:
				tag := combineTypeAndField(LengthDelim, protoField)
				err := writeTag(w, tag)
				if err != nil {
					return err
				}

				str := val.Interface().(string)
				err = writeLengthDelimited(w, []byte(str))
				if err != nil {
					return err
				}
			case reflect.Int32:
				tag := combineTypeAndField(Varint, protoField)
				err := writeTag(w, tag)
				if err != nil {
					return err
				}

				err = writeVarint(w, uint64(val.Interface().(int32)))
				if err != nil {
					return err
				}
			case reflect.Int64:
				tag := combineTypeAndField(Varint, protoField)
				err := writeTag(w, tag)
				if err != nil {
					return err
				}

				err = writeVarint(w, uint64(val.Interface().(int64)))
				if err != nil {
					return err
				}
			case reflect.Bool:
				b := val.Interface().(bool)
				var boolval uint64
				if b {
					boolval = 1
				}

				tag := combineTypeAndField(Varint, protoField)
				err := writeTag(w, tag)
				if err != nil {
					return err
				}

				err = writeVarint(w, boolval)
				if err != nil {
					return err
				}
			case reflect.Slice:
				if val.Type().Elem().Kind() != reflect.Uint8 {
					return fmt.Errorf("cannot handle arrays of non byte type")
				}
				tag := combineTypeAndField(LengthDelim, protoField)
				err := writeTag(w, tag)
				if err != nil {
					return err
				}

				byteval := val.Interface().([]byte)
				err = writeLengthDelimited(w, byteval)
				if err != nil {
					return err
				}

			default:
				return fmt.Errorf("unhandled type: %s", val)
			}
		}
	}

	for protoField, fprop := range props.FieldMapping {
		if fprop.Repeated {
			field := val.Field(fprop.GoField)
			if field.Kind() != reflect.Chan {
				return errors.New("repeated field was not a channel")
			}

			go handleChannelIn(w, sm, protoField, field.Type().Elem(), field)
		}
	}

	return nil
}

func handleChannelIn(w io.Writer, sm StreamMessage, field byte, typ reflect.Type, ch reflect.Value) {
	for {
		val, ok := ch.Recv()
		if !ok {
			return
		}

		switch val := val.Interface().(type) {
		case proto.Message:
			data, err := proto.Marshal(val)
			if err != nil {
				sm.Errors() <- err
				return
			}

			tag := combineTypeAndField(LengthDelim, field)
			err = writeTag(w, tag)
			if err != nil {
				sm.Errors() <- err
				return
			}

			err = writeLengthDelimited(w, data)
			if err != nil {
				sm.Errors() <- err
				return
			}
		case string:
			tag := combineTypeAndField(LengthDelim, field)
			err := writeTag(w, tag)
			if err != nil {
				sm.Errors() <- err
				return
			}

			err = writeLengthDelimited(w, []byte(val))
			if err != nil {
				sm.Errors() <- err
				return
			}
		case []byte:
			tag := combineTypeAndField(LengthDelim, field)
			err := writeTag(w, tag)
			if err != nil {
				sm.Errors() <- err
				return
			}

			err = writeLengthDelimited(w, val)
			if err != nil {
				sm.Errors() <- err
				return
			}
		case int32:
			tag := combineTypeAndField(Varint, field)
			err := writeTag(w, tag)
			if err != nil {
				sm.Errors() <- err
				return
			}

			err = writeVarint(w, uint64(val))
			if err != nil {
				sm.Errors() <- err
				return
			}
		default:
			fmt.Println("UNRECOGNIZED REPEATED FIELD TYPE", reflect.TypeOf(val))
			sm.Errors() <- errors.New("unrecognized repeated field type")
			return
		}

		_ = val
	}
}

type FieldInfo struct {
	// The field index in the Go struct
	GoField  int
	Repeated bool
	Type     string
}

type Props struct {
	// A mapping from the protobuf field number to field info
	FieldMapping map[byte]FieldInfo
}

func GetProperties(i proto.Message) (*Props, error) {
	t := reflect.TypeOf(i).Elem()

	props := &Props{FieldMapping: make(map[byte]FieldInfo)}
	for i := 0; i < t.NumField(); i++ {
		field := FieldInfo{GoField: i}

		tag := t.Field(i).Tag.Get("protobuf")
		if len(tag) == 0 {
			continue
		}

		parts := strings.Split(tag, ",")
		if len(parts) < 3 {
			return nil, errors.New("not enough values in protobuf field tag")
		}

		n, err := strconv.Atoi(parts[1])
		if err != nil {
			return nil, err
		}

		if parts[2] == "rep" {
			field.Repeated = true
		}

		field.Type = parts[0]
		props.FieldMapping[byte(n)] = field
	}
	return props, nil
}
<commit_msg>clean up encoding a bunch<commit_after>package pbs

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/golang/protobuf/proto"
	"io"
	"reflect"
	"strconv"
	"strings"
	"sync"
)

const (
	Varint      = 0
	Int64       = 1
	LengthDelim = 2
	StartGroup  = 3
	EndGroup    = 4
	Bit32       = 5
)

type StreamMessage interface {
	proto.Message
	io.Closer
	Errors() chan error
	Closed() <-chan struct{}
}

func splitTypeAndField(b byte) (Type, Field byte) {
	return (b & 0x7), b >> 3
}

func readLengthDelim(r *bufio.Reader) ([]byte, error) {
	l, err := readVarint(r)
	if err != nil {
		return nil, err
	}

	buf := make([]byte, l)
	_, err = io.ReadFull(r, buf)
	if err != nil {
		return nil, err
	}
	return buf, nil
}

func readVarint(r *bufio.Reader) (int, error) {
	var sum int
	for i := uint(0); i < 4; i++ {
		b, err := r.ReadByte()
		if err != nil {
			return 0, err
		}

		cont := b & 128
		val := b & 127

		sum += int(val << (7 * i))
		if cont == 0 {
			break
		}
	}
	return sum, nil
}

func (db *decBuffer) decodeField(field byte, data []byte) error {
	finfo := db.props.FieldMapping[field]

	val := reflect.ValueOf(db.val).Elem()

	f := val.Field(finfo.GoField)
	switch {
	case f.Type().Kind() == reflect.Chan:
		// This is a 'repeated' field
		// we just encode this value and send it along

		e := f.Type().Elem()
		if e.Kind() == reflect.Ptr {
			e = e.Elem()
		}

		newVal := reflect.New(e)
		switch pv := newVal.Interface().(type) {
		case proto.Message:
			err := proto.Unmarshal(data, pv)
			if err != nil {
				return err
			}
			f.Send(newVal)
		case *string:
			*pv = string(data)
			f.Send(newVal.Elem())
		case *[]byte:
			*pv = data
			f.Send(newVal.Elem())
		default:
			fmt.Println(reflect.TypeOf(pv))
			fmt.Println(data)
			return fmt.Errorf("Unrecognized type in protobuf field decode")
		}
	case f.Type().Elem().Kind() == reflect.String:
		f.Set(reflect.New(f.Type().Elem()))
		f.Elem().SetString(string(data))
	default:
		fmt.Println("UNKNOWN")
		fmt.Println(f)
		return errors.New("TODO: fix this unsupported type")
	}
	return nil
}

type decBuffer struct {
	props *Props
	val   StreamMessage
}

func StreamDecode(r io.Reader, sm StreamMessage) error {
	props, err := GetProperties(sm)
	if err != nil {
		return err
	}

	go func() {
		defer sm.Close()
		read := bufio.NewReader(r)

		db := decBuffer{
			props: props,
			val:   sm,
		}
		for {
			b, err := read.ReadByte()
			if err != nil {
				if err == io.EOF {
					return
				}
				sm.Errors() <- err
				return
			}

			typ, f := splitTypeAndField(b)
			switch typ {
			case Varint:
				i, err := readVarint(read)
				if err != nil {
					sm.Errors() <- err
					return
				}

				field := reflect.ValueOf(sm).Elem().Field(props.FieldMapping[f].GoField)

				fmt.Println("proto field: ", f)
				fmt.Println("FIELD: ", field)
				elemType := field.Type()
				if elemType.Kind() == reflect.Chan {
					elemType = elemType.Elem()
				}
				fmt.Println("elemtype: ", elemType)

				nval := reflect.New(elemType.Elem())
				nval.Elem().SetInt(int64(i))

				if props.FieldMapping[f].Repeated {
					field.Send(nval.Elem())
				} else {
					field.Set(nval)
				}
			case LengthDelim:
				val, err := readLengthDelim(read)
				if err != nil {
					if err == io.EOF {
						return
					}
					sm.Errors() <- err
					return
				}
				err = db.decodeField(f, val)
				if err != nil {
					sm.Errors() <- err
					return
				}
			default:
				fmt.Println("unrecognized type: ", typ)
				return
			}
		}
	}()
	return nil
}

func combineTypeAndField(typ, field byte) byte {
	return (field << 3) | (typ & 0x7)
}

func writeLengthDelimited(w io.Writer, field byte, data []byte) error {
	tag := combineTypeAndField(LengthDelim, field)
	err := writeTag(w, tag)
	if err != nil {
		return err
	}

	length := proto.EncodeVarint(uint64(len(data)))
	n, err := w.Write(length)
	if err != nil {
		return err
	}
	if n != len(length) {
		return errors.New("failed to write length")
	}

	n, err = w.Write([]byte(data))
	if err != nil {
		return err
	}
	if n != len(data) {
		return errors.New("failed to write data")
	}
	return nil
}

func writeTag(w io.Writer, tag byte) error {
	n, err := w.Write([]byte{tag})
	if err != nil {
		return err
	}
	if n != 1 {
		return errors.New("failed to write tag")
	}
	return nil
}

func writeVarint(w io.Writer, field byte, v uint64) error {
	tag := combineTypeAndField(Varint, field)
	err := writeTag(w, tag)
	if err != nil {
		return err
	}

	data := proto.EncodeVarint(v)
	n, err := w.Write(data)
	if err != nil {
		return err
	}
	if n != len(data) {
		return errors.New("failed to write enough bytes")
	}
	return nil
}

func writeProtoVal(w io.Writer, sm StreamMessage, field byte, val interface{}) error {
	switch val := val.(type) {
	case proto.Message:
		data, err := proto.Marshal(val)
		if err != nil {
			return err
		}

		err = writeLengthDelimited(w, field, data)
		if err != nil {
			return err
		}
	case string:
		err := writeLengthDelimited(w, field, []byte(val))
		if err != nil {
			return err
		}
	case []byte:
		err := writeLengthDelimited(w, field, val)
		if err != nil {
			return err
		}
	case int32:
		err := writeVarint(w, field, uint64(val))
		if err != nil {
			return err
		}
	case int64:
		err := writeVarint(w, field, uint64(val))
		if err != nil {
			return err
		}
	case bool:
		var boolval uint64
		if val {
			boolval = 1
		}

		err := writeVarint(w, field, boolval)
		if err != nil {
			return err
		}
	default:
		fmt.Println("UNRECOGNIZED REPEATED FIELD TYPE", reflect.TypeOf(val))
		return errors.New("unrecognized repeated field type")
	}

	return nil
}

// StreamEncode will perform a streaming encode of the given protobuf
// StreamMessage and write out the given writer. This function will
// return when all non-channel fields have been encoded, and goroutines
// will be spawned for the encoding of the channeled values. Those goroutines
// will receive on the channels and send values along as they get them until
// the StreamMessage is closed
func StreamEncode(w io.Writer, sm StreamMessage) error {
	// Parse out the protobuf struct tags
	props, err := GetProperties(sm)
	if err != nil {
		return err
	}

	val := reflect.ValueOf(sm).Elem()

	se := &streamEncoder{out: w, sm: sm}

	for protoField, fprop := range props.FieldMapping {
		field := val.Field(fprop.GoField)
		if fprop.Repeated {
			// Sanity check
			if field.Kind() != reflect.Chan {
				return errors.New("repeated field was not a channel")
			}

			go se.handleChannelIn(protoField, field)

		} else {
			if field.Kind() == reflect.Ptr {
				field = field.Elem()
			}
			err := writeProtoVal(w, sm, protoField, field.Interface())
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// streamEncoder is a helper struct to ensure that concurrent writes
// dont get intermingled.
type streamEncoder struct {
	out io.Writer
	sm  StreamMessage
	lk  sync.Mutex
}

func (se *streamEncoder) handleChannelIn(field byte, ch reflect.Value) {
	for {
		val, ok := ch.Recv()
		if !ok {
			return
		}

		se.lk.Lock()
		writeProtoVal(se.out, se.sm, field, val.Interface())
		se.lk.Unlock()
	}
}

type FieldInfo struct {
	// The field index in the Go struct
	GoField  int
	Repeated bool
	Type     string
}

type Props struct {
	// A mapping from the protobuf field number to field info
	FieldMapping map[byte]FieldInfo
}

func GetProperties(i proto.Message) (*Props, error) {
	t := reflect.TypeOf(i).Elem()

	props := &Props{FieldMapping: make(map[byte]FieldInfo)}
	for i := 0; i < t.NumField(); i++ {
		field := FieldInfo{GoField: i}

		tag := t.Field(i).Tag.Get("protobuf")
		if len(tag) == 0 {
			continue
		}

		parts := strings.Split(tag, ",")
		if len(parts) < 3 {
			return nil, errors.New("not enough values in protobuf field tag")
		}

		n, err := strconv.Atoi(parts[1])
		if err != nil {
			return nil, err
		}

		if parts[2] == "rep" {
			field.Repeated = true
		}

		field.Type = parts[0]
		props.FieldMapping[byte(n)] = field
	}
	return props, nil
}
<|endoftext|>


<commit_before>package streaming

import "io"
import "bufio"
import "github.com/tarm/goserial"
import "github.com/joushou/gocnc/gcode"
import "errors"
import "fmt"
import "time"

type Result struct {
	level   string
	message string
}

type Streamer struct {
	serialPort io.ReadWriteCloser
	reader     *bufio.Reader
	writer     *bufio.Writer
}

func serialReader(reader *bufio.Reader) Result {
	c, err := reader.ReadBytes('\n')
	if err != nil {
		return Result{"serial-error", fmt.Sprintf("%s", err)}
	}
	b := string(c)
	if b == "ok\r\n" {
		return Result{"ok", ""}
	} else if len(b) >= 5 && b[:5] == "error" {
		return Result{"error", b[6 : len(b)-1]}
	} else if len(b) >= 5 && b[:5] == "alarm" {
		return Result{"alarm", b[6 : len(b)-1]}
	} else {
		return Result{"info", b[:len(b)-1]}
	}
}

func (s *Streamer) Connect(name string) error {
	c := &serial.Config{Name: name, Baud: 115200}
	var err error
	s.serialPort, err = serial.OpenPort(c)
	if err != nil {
		return errors.New("Unable to connect to CNC!")
	}

	s.reader = bufio.NewReader(s.serialPort)
	s.writer = bufio.NewWriter(s.serialPort)

	for {
		c, err := s.reader.ReadBytes('\n')
		m := string(c)
		if len(m) == 26 && m[:5] == "Grbl " && m[9:] == " ['$' for help]\r\n" {
			fmt.Printf("Grbl version %s initialized\n", m[5:9])
			break
		} else if m == "\r\n" {
			continue
		}

		if err != nil {
			return errors.New("Unable to detect initialized GRBL")
		}
	}

	return nil
}

func (s *Streamer) Stop() {
	_, _ = s.serialPort.Write([]byte("\x18\n"))
	s.serialPort.Close()
}

func (s *Streamer) Send(doc *gcode.Document, maxPrecision int, progress chan int) error {
	defer close(progress)
	for idx, block := range doc.Blocks {
		_, err := s.writer.WriteString(block.Export(maxPrecision) + "\n")
		if err != nil {
			return errors.New("\nError while sending data:" + fmt.Sprintf("%s", err))
		}
		err = s.writer.Flush()
		if err != nil {
			return errors.New("\nError while flushing writer:" + fmt.Sprintf("%s", err))
		}

		res := serialReader(s.reader)

		switch res.level {
		case "error":
			return errors.New("Received error from CNC: " + res.message)
		case "alarm":
			return errors.New("Received alarm from CNC: " + res.message)
		case "info":
			fmt.Printf("\nReceived info from CNC: %s\n", res.message)
		}
		progress <- idx
		time.Sleep(1 * time.Microsecond)
	}
	return nil
}
<commit_msg>Use forked goserial<commit_after>package streaming

import "io"
import "bufio"
import "github.com/joushou/goserial"
import "github.com/joushou/gocnc/gcode"
import "errors"
import "fmt"

type Result struct {
	level   string
	message string
}

type Streamer struct {
	serialPort io.ReadWriteCloser
	reader     *bufio.Reader
	writer     *bufio.Writer
}

func serialReader(reader *bufio.Reader) Result {
	c, err := reader.ReadBytes('\n')
	if err != nil {
		return Result{"serial-error", fmt.Sprintf("%s", err)}
	}
	b := string(c)
	if b == "ok\r\n" {
		return Result{"ok", ""}
	} else if len(b) >= 5 && b[:5] == "error" {
		return Result{"error", b[6 : len(b)-1]}
	} else if len(b) >= 5 && b[:5] == "alarm" {
		return Result{"alarm", b[6 : len(b)-1]}
	} else {
		return Result{"info", b[:len(b)-1]}
	}
}

func (s *Streamer) Connect(name string) error {
	c := &serial.Config{Name: name, Baud: 115200}
	var err error
	s.serialPort, err = serial.OpenPort(c)
	if err != nil {
		return err
	}

	s.reader = bufio.NewReader(s.serialPort)
	s.writer = bufio.NewWriter(s.serialPort)

	for {
		c, err := s.reader.ReadBytes('\n')
		m := string(c)
		if len(m) == 26 && m[:5] == "Grbl " && m[9:] == " ['$' for help]\r\n" {
			fmt.Printf("Grbl version %s initialized\n", m[5:9])
			break
		} else if m == "\r\n" {
			continue
		}

		if err != nil {
			return errors.New("Unable to detect initialized GRBL")
		}
	}

	return nil
}

func (s *Streamer) Stop() {
	_, _ = s.serialPort.Write([]byte("\x18\n"))
	s.serialPort.Close()
}

func (s *Streamer) Send(doc *gcode.Document, maxPrecision int, progress chan int) error {
	defer close(progress)
	for idx, block := range doc.Blocks {
		_, err := s.writer.WriteString(block.Export(maxPrecision) + "\n")
		if err != nil {
			return errors.New("\nError while sending data:" + fmt.Sprintf("%s", err))
		}
		err = s.writer.Flush()
		if err != nil {
			return errors.New("\nError while flushing writer:" + fmt.Sprintf("%s", err))
		}

		res := serialReader(s.reader)

		switch res.level {
		case "error":
			return errors.New("Received error from CNC: " + res.message)
		case "alarm":
			return errors.New("Received alarm from CNC: " + res.message)
		case "info":
			fmt.Printf("\nReceived info from CNC: %s\n", res.message)
		}
		progress <- idx
	}
	return nil
}
<|endoftext|>