<commit_before>// Copyright 2018 The Wuffs Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://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.

// +build ignore

package main

// This program exercises the Go standard library's GIF decoder.
//
// Wuffs' C code doesn't depend on Go per se, but this program gives some
// performance data for specific Go GIF implementations. The equivalent Wuffs
// benchmarks (on the same test images) are run via:
//
// wuffs bench std/gif

import (
	"bytes"
	"fmt"
	"image"
	"image/draw"
	"image/gif"
	"io/ioutil"
	"os"
	"runtime"
	"strings"
	"time"
)

const (
	iterscale = 20
	reps      = 5
)

type testCase = struct {
	benchname     string
	filename      string
	itersUnscaled uint32
}

var testCases = []testCase{{
	benchname:     "go_gif_decode_1k_bw",
	filename:      "test/data/pjw-thumbnail.gif",
	itersUnscaled: 2000,
}, {
	benchname:     "go_gif_decode_1k_color",
	filename:      "test/data/hippopotamus.regular.gif",
	itersUnscaled: 1000,
}, {
	benchname:     "go_gif_decode_10k_bgra",
	filename:      "test/data/hat.gif",
	itersUnscaled: 100,
}, {
	benchname:     "go_gif_decode_10k_indexed",
	filename:      "test/data/hat.gif",
	itersUnscaled: 100,
}, {
	benchname:     "go_gif_decode_20k",
	filename:      "test/data/bricks-gray.gif",
	itersUnscaled: 50,
}, {
	benchname:     "go_gif_decode_100k_artificial",
	filename:      "test/data/hibiscus.primitive.gif",
	itersUnscaled: 15,
}, {
	benchname:     "go_gif_decode_100k_realistic",
	filename:      "test/data/hibiscus.regular.gif",
	itersUnscaled: 10,
}, {
	benchname:     "go_gif_decode_1000k",
	filename:      "test/data/harvesters.gif",
	itersUnscaled: 1,
}, {
	benchname:     "go_gif_decode_anim_screencap",
	filename:      "test/data/gifplayer-muybridge.gif",
	itersUnscaled: 1,
}}

func main() {
	if err := main1(); err != nil {
		os.Stderr.WriteString(err.Error() + "\n")
		os.Exit(1)
	}
}

func main1() error {
	fmt.Printf("# Go %s\n", runtime.Version())
	fmt.Printf("#\n")
	fmt.Printf("# The output format, including the \"Benchmark\" prefixes, is compatible with the\n")
	fmt.Printf("# https://godoc.org/golang.org/x/perf/cmd/benchstat tool. To install it, first\n")
	fmt.Printf("# install Go, then run \"go get golang.org/x/perf/cmd/benchstat\".\n")

	for i := -1; i < reps; i++ {
		for _, tc := range testCases {
			src, err := ioutil.ReadFile("../../" + tc.filename)
			if err != nil {
				return err
			}

			runtime.GC()

			start := time.Now()

			iters := uint64(tc.itersUnscaled) * iterscale
			bgra := strings.HasSuffix(tc.benchname, "_bgra")
			numBytes, err := decode(src, bgra)
			if err != nil {
				return err
			}
			for j := uint64(1); j < iters; j++ {
				decode(src, bgra)
			}

			elapsedNanos := time.Since(start)

			kbPerS := numBytes * uint64(iters) * 1000000 / uint64(elapsedNanos)

			if i < 0 {
				continue // Warm up rep.
			}

			fmt.Printf("Benchmark%-30s %8d %12d ns/op %8d.%03d MB/s\n",
				tc.benchname, iters, uint64(elapsedNanos)/iters, kbPerS/1000, kbPerS%1000)
		}
	}

	return nil
}

func decode(src []byte, bgra bool) (numBytes uint64, retErr error) {
	g, err := gif.DecodeAll(bytes.NewReader(src))
	if err != nil {
		return 0, err
	}

	for _, m := range g.Image {
		b := m.Bounds()
		n := uint64(b.Dx()) * uint64(b.Dy())

		// Go converts to RGBA (as that's what Go's image/draw standard library is
		// optimized for); Wuffs converts to BGRA. The difference isn't important,
		// as we just want measure how long it takes.
		if bgra {
			dst := image.NewRGBA(b)
			draw.Draw(dst, b, m, b.Min, draw.Src)
			n *= 4
		}

		numBytes += n
	}

	return numBytes, nil
}
<commit_msg>Have script/bench-go-gif load files less often<commit_after>// Copyright 2018 The Wuffs Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://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.

// +build ignore

package main

// This program exercises the Go standard library's GIF decoder.
//
// Wuffs' C code doesn't depend on Go per se, but this program gives some
// performance data for specific Go GIF implementations. The equivalent Wuffs
// benchmarks (on the same test images) are run via:
//
// wuffs bench std/gif

import (
	"bytes"
	"fmt"
	"image"
	"image/draw"
	"image/gif"
	"io/ioutil"
	"os"
	"runtime"
	"strings"
	"time"
)

const (
	iterscale = 20
	reps      = 5
)

type testCase = struct {
	benchname     string
	src           []byte
	itersUnscaled uint32
}

var testCases = []testCase{{
	benchname:     "go_gif_decode_1k_bw",
	src:           mustLoad("test/data/pjw-thumbnail.gif"),
	itersUnscaled: 2000,
}, {
	benchname:     "go_gif_decode_1k_color",
	src:           mustLoad("test/data/hippopotamus.regular.gif"),
	itersUnscaled: 1000,
}, {
	benchname:     "go_gif_decode_10k_bgra",
	src:           mustLoad("test/data/hat.gif"),
	itersUnscaled: 100,
}, {
	benchname:     "go_gif_decode_10k_indexed",
	src:           mustLoad("test/data/hat.gif"),
	itersUnscaled: 100,
}, {
	benchname:     "go_gif_decode_20k",
	src:           mustLoad("test/data/bricks-gray.gif"),
	itersUnscaled: 50,
}, {
	benchname:     "go_gif_decode_100k_artificial",
	src:           mustLoad("test/data/hibiscus.primitive.gif"),
	itersUnscaled: 15,
}, {
	benchname:     "go_gif_decode_100k_realistic",
	src:           mustLoad("test/data/hibiscus.regular.gif"),
	itersUnscaled: 10,
}, {
	benchname:     "go_gif_decode_1000k",
	src:           mustLoad("test/data/harvesters.gif"),
	itersUnscaled: 1,
}, {
	benchname:     "go_gif_decode_anim_screencap",
	src:           mustLoad("test/data/gifplayer-muybridge.gif"),
	itersUnscaled: 1,
}}

func mustLoad(filename string) []byte {
	src, err := ioutil.ReadFile("../../" + filename)
	if err != nil {
		panic(err.Error())
	}
	return src
}

func main() {
	if err := main1(); err != nil {
		os.Stderr.WriteString(err.Error() + "\n")
		os.Exit(1)
	}
}

func main1() error {
	fmt.Printf("# Go %s\n", runtime.Version())
	fmt.Printf("#\n")
	fmt.Printf("# The output format, including the \"Benchmark\" prefixes, is compatible with the\n")
	fmt.Printf("# https://godoc.org/golang.org/x/perf/cmd/benchstat tool. To install it, first\n")
	fmt.Printf("# install Go, then run \"go get golang.org/x/perf/cmd/benchstat\".\n")

	for i := -1; i < reps; i++ {
		for _, tc := range testCases {
			runtime.GC()

			start := time.Now()

			iters := uint64(tc.itersUnscaled) * iterscale
			bgra := strings.HasSuffix(tc.benchname, "_bgra")
			numBytes, err := decode(tc.src, bgra)
			if err != nil {
				return err
			}
			for j := uint64(1); j < iters; j++ {
				decode(tc.src, bgra)
			}

			elapsedNanos := time.Since(start)

			kbPerS := numBytes * uint64(iters) * 1000000 / uint64(elapsedNanos)

			if i < 0 {
				continue // Warm up rep.
			}

			fmt.Printf("Benchmark%-30s %8d %12d ns/op %8d.%03d MB/s\n",
				tc.benchname, iters, uint64(elapsedNanos)/iters, kbPerS/1000, kbPerS%1000)
		}
	}

	return nil
}

func decode(src []byte, bgra bool) (numBytes uint64, retErr error) {
	g, err := gif.DecodeAll(bytes.NewReader(src))
	if err != nil {
		return 0, err
	}

	for _, m := range g.Image {
		b := m.Bounds()
		n := uint64(b.Dx()) * uint64(b.Dy())

		// Go converts to RGBA (as that's what Go's image/draw standard library is
		// optimized for); Wuffs converts to BGRA. The difference isn't important,
		// as we just want measure how long it takes.
		if bgra {
			dst := image.NewRGBA(b)
			draw.Draw(dst, b, m, b.Min, draw.Src)
			n *= 4
		}

		numBytes += n
	}

	return numBytes, nil
}
<|endoftext|>


<commit_before>package widget

import (
	"strings"

	"github.com/burl/termbox-go"
)

// StrBuf is the basis for an input field widget
type StrBuf struct {
	Buf       string
	Pos       int
	Row       int
	Col       int
	isMasked  bool
	inputMask rune
}

// NewStrBuf - create a new StrBuf at x,y in the screen
func NewStrBuf(x, y int) *StrBuf {
	return &StrBuf{
		Buf: "",
		Pos: 0,
		Row: y,
		Col: x,
	}
}

// MaskInput - mask input (echo a mask char instead of actual char)
func (b *StrBuf) MaskInput(ch rune) {
	b.isMasked = true
	b.inputMask = ch
}

// Draw the buffer (not so performant.. could be optimized) - and draw
// a "cursor" by showing inverse...
//
// I previously turned on the real terminal cursor and positioned it, but
// there was a lot of extra accounting that had to take place.  This just
// draws an inverse character (or space) at the insertion point, it looks
// fine on most terminals and works with white backgrounds, etc.
//
func (b *StrBuf) Draw() {
	buf := b.Buf
	if b.isMasked {
		buf = strings.Repeat(string(b.inputMask), len(buf))
	}
	tbPrint(b.Col+1, b.Row, coldef, coldef, buf+"\x20\x20")
	ch := '\x20'
	if b.Pos < len(b.Buf) {
		if b.isMasked {
			ch = '*'
		} else {
			ch = rune(b.Buf[b.Pos])
		}
	}
	termbox.SetCell(b.Pos+b.Col+1, 0, ch, termbox.AttrReverse, coldef)
}

// SetValue set the value of the StrBuf internal string
func (b *StrBuf) SetValue(value string) {
	for i := 0; i < len(value); i++ {
		b.Insert(rune(value[i]))
	}
}

// String - return value of strbuf
func (b *StrBuf) String() string {
	return b.Buf
}

// Append - add a character to the end of the StrBuf
func (b *StrBuf) Append(ch rune) {
	b.Buf = b.Buf + string(ch)
	b.Right()
}

// Insert a character at the current edit point in the StrBuf
func (b *StrBuf) Insert(ch rune) {
	b.Buf = b.Buf[0:b.Pos] + string(ch) + b.Buf[b.Pos:]
	b.Right()
}

// Delete a character at the current position in the StrBuf
func (b *StrBuf) Delete() {
	if len(b.Buf) < 1 {
		return
	}
	if b.Pos == 0 {
		return
	}
	b.Buf = b.Buf[0:b.Pos-1] + b.Buf[b.Pos:]
	b.Left()
}

// Left - move cursor left in the StrBuf
func (b *StrBuf) Left() {
	if b.Pos > 0 {
		b.Pos = b.Pos - 1
		b.Draw()
	}
}

// Beginning - move cursor/insertion point to beginning of StrBuf
func (b *StrBuf) Beginning() {
	if b.Pos > 0 {
		b.Pos = 0
		b.Draw()
	}
}

// Right - move cursor/insertion-point to the right in the StrBuf
func (b *StrBuf) Right() {
	if b.Pos < len(b.Buf) {
		b.Pos = b.Pos + 1
		b.Draw()
	}
}

// End - move cursor/insertion point to the end of the StrBuf
func (b *StrBuf) End() {
	dx := len(b.Buf) - b.Pos
	if dx > 1 {
		b.Pos += dx
		b.Draw()
	}
}
<commit_msg>fixing mistake with mask char<commit_after>package widget

import (
	"strings"

	"github.com/burl/termbox-go"
)

// StrBuf is the basis for an input field widget
type StrBuf struct {
	Buf       string
	Pos       int
	Row       int
	Col       int
	isMasked  bool
	inputMask rune
}

// NewStrBuf - create a new StrBuf at x,y in the screen
func NewStrBuf(x, y int) *StrBuf {
	return &StrBuf{
		Buf: "",
		Pos: 0,
		Row: y,
		Col: x,
	}
}

// MaskInput - mask input (echo a mask char instead of actual char)
func (b *StrBuf) MaskInput(ch rune) {
	b.isMasked = true
	b.inputMask = ch
}

// Draw the buffer (not so performant.. could be optimized) - and draw
// a "cursor" by showing inverse...
//
// I previously turned on the real terminal cursor and positioned it, but
// there was a lot of extra accounting that had to take place.  This just
// draws an inverse character (or space) at the insertion point, it looks
// fine on most terminals and works with white backgrounds, etc.
//
func (b *StrBuf) Draw() {
	buf := b.Buf
	if b.isMasked {
		buf = strings.Repeat(string(b.inputMask), len(buf))
	}
	tbPrint(b.Col+1, b.Row, coldef, coldef, buf+"\x20\x20")
	ch := '\x20'
	if b.Pos < len(b.Buf) {
		if b.isMasked {
			ch = b.inputMask
		} else {
			ch = rune(b.Buf[b.Pos])
		}
	}
	termbox.SetCell(b.Pos+b.Col+1, 0, ch, termbox.AttrReverse, coldef)
}

// SetValue set the value of the StrBuf internal string
func (b *StrBuf) SetValue(value string) {
	for i := 0; i < len(value); i++ {
		b.Insert(rune(value[i]))
	}
}

// String - return value of strbuf
func (b *StrBuf) String() string {
	return b.Buf
}

// Append - add a character to the end of the StrBuf
func (b *StrBuf) Append(ch rune) {
	b.Buf = b.Buf + string(ch)
	b.Right()
}

// Insert a character at the current edit point in the StrBuf
func (b *StrBuf) Insert(ch rune) {
	b.Buf = b.Buf[0:b.Pos] + string(ch) + b.Buf[b.Pos:]
	b.Right()
}

// Delete a character at the current position in the StrBuf
func (b *StrBuf) Delete() {
	if len(b.Buf) < 1 {
		return
	}
	if b.Pos == 0 {
		return
	}
	b.Buf = b.Buf[0:b.Pos-1] + b.Buf[b.Pos:]
	b.Left()
}

// Left - move cursor left in the StrBuf
func (b *StrBuf) Left() {
	if b.Pos > 0 {
		b.Pos = b.Pos - 1
		b.Draw()
	}
}

// Beginning - move cursor/insertion point to beginning of StrBuf
func (b *StrBuf) Beginning() {
	if b.Pos > 0 {
		b.Pos = 0
		b.Draw()
	}
}

// Right - move cursor/insertion-point to the right in the StrBuf
func (b *StrBuf) Right() {
	if b.Pos < len(b.Buf) {
		b.Pos = b.Pos + 1
		b.Draw()
	}
}

// End - move cursor/insertion point to the end of the StrBuf
func (b *StrBuf) End() {
	dx := len(b.Buf) - b.Pos
	if dx > 1 {
		b.Pos += dx
		b.Draw()
	}
}
<|endoftext|>


<commit_before>package worker

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"time"

	"github.com/peterstace/grayt/protocol"
	"github.com/peterstace/grayt/trace"
)

func NewServer(scenelibAddr string) *Server {
	return &Server{
		scenelibAddr: scenelibAddr,
		scenes:       map[string]trace.Scene{},
	}
}

type Server struct {
	scenelibAddr string

	mu      sync.Mutex
	working bool

	scenes map[string]trace.Scene
}

func (s *Server) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	if req.URL.Path == "/trace" {
		s.handleTrace(w, req)
		return
	}
	http.NotFound(w, req)
}

func (s *Server) handleTrace(w http.ResponseWriter, req *http.Request) {
	if req.Method != http.MethodGet {
		http.Error(w, "only GET allowed", http.StatusMethodNotAllowed)
		return
	}
	params := req.URL.Query()
	sceneName := params.Get("scene_name")
	if sceneName == "" {
		http.Error(w, "scene_name query param not set", http.StatusBadRequest)
		return
	}
	pxWideStr := params.Get("px_wide")
	if pxWideStr == "" {
		http.Error(w, "px_wide query param not set", http.StatusBadRequest)
		return
	}
	pxWide, err := strconv.Atoi(pxWideStr)
	if err != nil {
		http.Error(w, "couldn't convert px_wide to int", http.StatusBadRequest)
		return
	}
	pxHighStr := params.Get("px_high")
	if pxHighStr == "" {
		http.Error(w, "px_high query param not set", http.StatusBadRequest)
		return
	}
	pxHigh, err := strconv.Atoi(pxHighStr)
	if err != nil {
		http.Error(w, "couldn't convert px_high to int", http.StatusBadRequest)
		return
	}
	s.serveLayer(w, sceneName, pxWide, pxHigh)
}

func (s *Server) serveLayer(
	w http.ResponseWriter, sceneName string, pxWide, pxHigh int,
) {
	s.mu.Lock()
	if s.working {
		http.Error(w, "already working", http.StatusTooManyRequests)
		s.mu.Unlock()
		return
	}
	s.working = true
	s.mu.Unlock()

	defer func() {
		s.mu.Lock()
		s.working = false
		s.mu.Unlock()
	}()

	scene, ok := s.scenes[sceneName]
	if !ok {
		log.Printf("scene %q not cached, fetching from scenelib", sceneName)
		resp, err := http.Get(
			s.scenelibAddr + "/scene?name=" + url.PathEscape(sceneName),
		)
		if err != nil {
			http.Error(w,
				"fetching scene: "+err.Error(),
				http.StatusInternalServerError,
			)
			return
		}
		if resp.StatusCode != http.StatusOK {
			w.WriteHeader(resp.StatusCode)
			fmt.Fprintf(w, "fetching scene: ")
			io.Copy(w, resp.Body)
			return
		}
		var sceneProto protocol.Scene
		if err := json.NewDecoder(resp.Body).Decode(&sceneProto); err != nil {
			http.Error(w,
				fmt.Sprintf("decoding scene: %v", err),
				http.StatusInternalServerError,
			)
			return
		}
		scene = trace.BuildScene(sceneProto)
		s.scenes[sceneName] = scene
	}

	// TODO Trace the image and return the data.

	time.Sleep(time.Second)
	fmt.Fprintf(w, "DATA GETS SEND HERE")
}
<commit_msg>Trace out a single layer<commit_after>package worker

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"time"

	"github.com/peterstace/grayt/protocol"
	"github.com/peterstace/grayt/trace"
)

func NewServer(scenelibAddr string) *Server {
	return &Server{
		scenelibAddr: scenelibAddr,
		scenes:       map[string]trace.Scene{},
	}
}

type Server struct {
	scenelibAddr string

	mu      sync.Mutex
	working bool

	scenes map[string]trace.Scene
}

func (s *Server) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	if req.URL.Path == "/trace" {
		s.handleTrace(w, req)
		return
	}
	http.NotFound(w, req)
}

func (s *Server) handleTrace(w http.ResponseWriter, req *http.Request) {
	if req.Method != http.MethodGet {
		http.Error(w, "only GET allowed", http.StatusMethodNotAllowed)
		return
	}
	params := req.URL.Query()
	sceneName := params.Get("scene_name")
	if sceneName == "" {
		http.Error(w, "scene_name query param not set", http.StatusBadRequest)
		return
	}
	pxWideStr := params.Get("px_wide")
	if pxWideStr == "" {
		http.Error(w, "px_wide query param not set", http.StatusBadRequest)
		return
	}
	pxWide, err := strconv.Atoi(pxWideStr)
	if err != nil {
		http.Error(w, "couldn't convert px_wide to int", http.StatusBadRequest)
		return
	}
	pxHighStr := params.Get("px_high")
	if pxHighStr == "" {
		http.Error(w, "px_high query param not set", http.StatusBadRequest)
		return
	}
	pxHigh, err := strconv.Atoi(pxHighStr)
	if err != nil {
		http.Error(w, "couldn't convert px_high to int", http.StatusBadRequest)
		return
	}
	s.serveLayer(w, sceneName, pxWide, pxHigh)
}

func (s *Server) serveLayer(
	w http.ResponseWriter, sceneName string, pxWide, pxHigh int,
) {
	s.mu.Lock()
	if s.working {
		http.Error(w, "already working", http.StatusTooManyRequests)
		s.mu.Unlock()
		return
	}
	s.working = true
	s.mu.Unlock()

	defer func() {
		s.mu.Lock()
		s.working = false
		s.mu.Unlock()
	}()

	scene, ok := s.scenes[sceneName]
	if !ok {
		log.Printf("scene %q not cached, fetching from scenelib", sceneName)
		resp, err := http.Get(
			s.scenelibAddr + "/scene?name=" + url.PathEscape(sceneName),
		)
		if err != nil {
			http.Error(w,
				"fetching scene: "+err.Error(),
				http.StatusInternalServerError,
			)
			return
		}
		if resp.StatusCode != http.StatusOK {
			w.WriteHeader(resp.StatusCode)
			fmt.Fprintf(w, "fetching scene: ")
			io.Copy(w, resp.Body)
			return
		}
		var sceneProto protocol.Scene
		if err := json.NewDecoder(resp.Body).Decode(&sceneProto); err != nil {
			http.Error(w,
				fmt.Sprintf("decoding scene: %v", err),
				http.StatusInternalServerError,
			)
			return
		}
		scene = trace.BuildScene(sceneProto)
		s.scenes[sceneName] = scene
	}

	traceLayer(w, pxWide, pxHigh, scene)
}

func traceLayer(w io.Writer, pxWide, pxHigh int, scene trace.Scene) {
	// TODO: a lot of this can be cached
	accel := trace.NewGrid(4, scene.Objects)
	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
	tr := trace.NewTracer(accel, rng)
	pxPitch := 2.0 / float64(pxWide)
	for pxY := 0; pxY < pxHigh; pxY++ {
		for pxX := 0; pxX < pxWide; pxX++ {
			x := (float64(pxX-pxWide/2) + rng.Float64()) * pxPitch
			y := (float64(pxY-pxHigh/2) + rng.Float64()) * pxPitch * -1.0
			cr := scene.Camera.MakeRay(x, y, rng)
			cr.Dir = cr.Dir.Unit()
			c := tr.TracePath(cr)
			binary.Write(w, binary.BigEndian, c.R)
			binary.Write(w, binary.BigEndian, c.G)
			binary.Write(w, binary.BigEndian, c.B)
		}
	}
}
<|endoftext|>


<commit_before>package wsproxy

import (
	"bufio"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/gorilla/websocket"
	"github.com/taskcluster/websocktunnel/util"
	"github.com/taskcluster/websocktunnel/wsmux"

	"github.com/sirupsen/logrus"
	nullLog "github.com/sirupsen/logrus/hooks/test"
)

const (
	monthUnix = 31 * 24 * time.Hour
)

var (
	clientIdRe = regexp.MustCompile(`^[a-zA-Z0-9_~.%-]+$`)
)

// Config contains the run time parameters for the proxy
type Config struct {
	// Upgrader is a websocket.Upgrader instance which is used to upgrade incoming
	// websocket connections from Clients.
	Upgrader websocket.Upgrader

	// Logger is used to log proxy events. Refer util.Logger.
	Logger *logrus.Logger

	// JWTSecretA and JWTSecretB are used by the proxy to verify JWTs from Clients.
	JWTSecretA []byte
	JWTSecretB []byte

	// Domain and port where proxy will be hosted
	Domain string
	Port   int

	// set to true if serving with TLS
	TLS bool

	// Audience value for aud claim
	Audience string
}

// proxy is used to send http and ws requests to a registered client.
// New proxy can be created by using wsproxy.New()
type proxy struct {
	m               sync.RWMutex
	pool            map[string]*wsmux.Session
	upgrader        websocket.Upgrader
	logger          *logrus.Logger
	onSessionRemove func(string)
	jwtSecretA      []byte
	jwtSecretB      []byte
	domain          string
	port            int
	tls             bool
	audience        string
}

// New creates a new proxy instance and wraps it as an http.Handler.
func New(conf Config) (http.Handler, error) {
	return newProxy(conf)
}

// ServeHTTP implements http.Handler so that the proxy may be used as a handler in a Mux or http.Server
func (p *proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	p.logf("", r.RemoteAddr, "Host=%s Path=%s", r.Host, r.URL.Path)

	// Client registration requests are a GET of path / with some headers set
	if path, id := r.URL.Path, r.Header.Get("x-websocktunnel-id"); id != "" && path == "/" {
		tokenString := util.ExtractJWT(r.Header.Get("Authorization"))
		p.register(w, r, id, tokenString)
		return
	}

	// try to match viewer requests to https://<domain>/<id>/<path> (ignoring
	// the host field)
	s := strings.TrimPrefix(r.URL.RequestURI(), "/")
	index := strings.Index(s, "/")
	id, path := "", "/"
	if index < 0 {
		id = s
	} else {
		id = s[:index]
		path = s[index:]
	}
	if id == "" {
		http.NotFound(w, r)
		return
	}
	p.serveRequest(w, r, id, path)
}

func newProxy(conf Config) (*proxy, error) {
	p := &proxy{
		pool:       make(map[string]*wsmux.Session),
		upgrader:   conf.Upgrader,
		logger:     conf.Logger,
		jwtSecretA: conf.JWTSecretA,
		jwtSecretB: conf.JWTSecretB,
		domain:     conf.Domain,
		port:       conf.Port,
		tls:        conf.TLS,
		audience:   conf.Audience,
	}

	if conf.Port == 0 {
		panic("no port specified")
	}

	if len(p.jwtSecretA) == 0 || len(p.jwtSecretB) == 0 {
		panic("wsproxy: missing secrets")
	}

	if p.logger == nil {
		logger, _ := nullLog.NewNullLogger()
		p.logger = logger
	}

	return p, nil

}

// setSessionRemoveHandler sets a function which is called when a wsmux Session is removed from
// the proxy due to closure or error.
func (p *proxy) setSessionRemoveHandler(h func(string)) {
	p.m.Lock()
	defer p.m.Unlock()
	p.onSessionRemove = h
}

// getWorkerSession returns true if a session with the given id is present
func (p *proxy) getWorkerSession(id string) (*wsmux.Session, bool) {
	p.m.RLock()
	defer p.m.RUnlock()
	s, ok := p.pool[id]
	return s, ok
}

// removeTunnel is an idempotent operation which deletes a client session from the proxy's
// pool
func (p *proxy) removeTunnel(id string) {
	p.m.Lock()
	defer p.m.Unlock()
	delete(p.pool, id)
	p.logf(id, "", "session removed")
}

// register is used to connect a client to the proxy so that it can start serving API endpoints.
// The request must contain the tunnel ID in the url.
// The request is validated by the proxy and the http connection is upgraded to websocket.
func (p *proxy) register(w http.ResponseWriter, r *http.Request, id, tokenString string) {
	if !clientIdRe.MatchString(id) {
		p.logerrorf(id, r.RemoteAddr, "client ID is invalid")
		http.Error(w, http.StatusText(400), 400)
		return
	}

	p.logf(id, r.RemoteAddr, "requesting client registration")
	if tokenString == "" {
		// No jwt. Connection not authorized
		p.logerrorf(id, r.RemoteAddr, "could not retreive auth token")
		http.Error(w, http.StatusText(400), 400)
		return
	}

	if !websocket.IsWebSocketUpgrade(r) {
		p.logerrorf(id, r.RemoteAddr, "request must be websocket upgrade")
		http.NotFound(w, r)
		return
	}

	// validation does not require lock
	if err := p.validateJWT(id, tokenString); err != nil {
		p.logerrorf(id, r.RemoteAddr, "unable to validate token: %v", err)
		http.Error(w, http.StatusText(401), 401)
		return
	}

	p.m.Lock()

	// remove any existing session forcibly
	for {
		existingSession := p.pool[id]
		if existingSession == nil {
			break
		}

		// If there's an existing session, unlock, close the existing session
		// (which will remove it), and try again.  Note that Session.Close is
		// properly reentrant, so two goroutines calling this at the same time
		// will not cause an issue.
		p.m.Unlock()
		_ = existingSession.Close()
		p.m.Lock()
	}

	defer p.m.Unlock()

	delete(p.pool, id)

	header := make(http.Header)

	urlScheme := "http://"
	defaultPort := 80
	if p.tls {
		urlScheme = "https://"
		defaultPort = 443
	}
	var portStr string
	if p.port != defaultPort {
		portStr = fmt.Sprintf(":%d", p.port)
	}

	url := urlScheme + p.domain + portStr + "/" + id
	header.Set("x-websocktunnel-client-url", url)
	p.logf(id, r.RemoteAddr, "sending url= %s", url)
	conn, err := p.upgrader.Upgrade(w, r, header)
	if err != nil {
		p.logger.Print(err)
		return
	}

	// generate config
	conf := wsmux.Config{
		StreamBufferSize: 4 * 1024,
		CloseCallback: func() {
			p.removeTunnel(id)
			if p.onSessionRemove != nil {
				p.onSessionRemove(id)
			}
		},
		Log: p.logger,
	}

	p.pool[id] = wsmux.Server(conn, conf)
	p.logf(id, r.RemoteAddr, "added new tunnel")
}

// serveRequest serves tunnel endpoints to viewers
func (p *proxy) serveRequest(w http.ResponseWriter, r *http.Request, id string, path string) {
	// log new request arrival
	p.logf(id, r.RemoteAddr, "request: host=%s path=%s", r.Host, path)
	p.logf(id, r.RemoteAddr, "request: URL: %v", r.URL)

	session, ok := p.getWorkerSession(id)

	// return 504 (bad gateway) if tunnel is not registered on this proxy
	if !ok {
		p.logerrorf(id, r.RemoteAddr, "could not find requested tunnel")
		http.Error(w, fmt.Sprintf("No client is connected with that id"), 504)
		return
	}

	// set original path as header
	r.Header.Set("x-websocktunnel-original-path", r.URL.Path)

	// check for a websocket request
	if websocket.IsWebSocketUpgrade(r) {
		_ = p.websocketProxy(w, r, session, id)
		return
	}

	// Open a stream to the tunnel session
	// path is the modified RequestURI
	reqURI, err := url.ParseRequestURI(path)
	if err != nil {
		http.Error(w, http.StatusText(500), 500)
		return
	}
	r.URL = reqURI
	p.logf(id, r.RemoteAddr, "attempting to open new stream")
	reqStream, err := session.Open()
	if err != nil {
		p.logerrorf(id, r.RemoteAddr, "could not open stream: path=%s", path)
		http.Error(w, http.StatusText(500), 500)
		return
	}

	// rewrite path for tunnel and write request
	err = r.Write(reqStream)
	if err != nil {
		p.logerrorf(id, r.RemoteAddr, "could not write request: path=%s", path)
		http.Error(w, http.StatusText(500), 500)
		return
	}

	// read response from tunnel
	bufReader := bufio.NewReader(reqStream)
	resp, err := http.ReadResponse(bufReader, r)
	if err != nil {
		p.logerrorf(id, r.RemoteAddr, "could not read response: path=%s", path)
		http.Error(w, http.StatusText(500), 500)
		return
	}

	// manually proxy response
	// clear responseWriter headers and write response headers instead
	for k := range w.Header() {
		w.Header().Del(k)
	}
	for k, v := range resp.Header {
		w.Header()[k] = v
	}

	// dump headers
	w.WriteHeader(resp.StatusCode)
	if resp.Body == nil {
		return
	}

	// stream body to viewer and close wsmux stream
	defer func() {
		_ = reqStream.Close()
	}()

	flusher, ok := w.(http.Flusher)
	// flusher may not be implemented by a ResponseWriter wrapper
	// simple copy
	if !ok {
		n, err := io.Copy(w, resp.Body)
		p.logf(id, r.RemoteAddr, "data transfered over request: %d bytes, error: %v", n, err)
		// log here
		return
	}

	p.logf(id, r.RemoteAddr, "streaming http")
	wf := &threadSafeWriteFlusher{w: w, f: flusher}
	n, err := copyAndFlush(wf, resp.Body, 100*time.Millisecond)
	p.logf(id, r.RemoteAddr, "data transfered over request: %d bytes, error: %v", n, err)
}

// validate jwt
// jwt signing and verification algorithm must be HMAC
func (p *proxy) validateJWT(id string, tokenString string) error {
	// parse jwt token
	// default parser verifies iat token if present. This can be a problem because of clocks not being
	// in sync.
	parser := &jwt.Parser{
		ValidMethods:         []string{"HS256"},
		SkipClaimsValidation: true, // Claims will be verified if token can be decoded using secret
	}

	token, err := parser.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, ErrUnexpectedSigningMethod
		}
		return p.jwtSecretA, nil
	})

	if err != nil {
		// log first error
		p.logerrorf(id, "", "%v: trying with second secret", err)

		token, err = parser.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
			if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
				return nil, ErrUnexpectedSigningMethod
			}
			return p.jwtSecretB, nil
		})
	}

	if err != nil {
		p.logerrorf(id, "", "%v: auth failed", err)
		return ErrAuthFailed
	}

	// check claims
	now := time.Now().Unix()
	claims, ok := token.Claims.(jwt.MapClaims)

	if !ok {
		p.logerrorf(id, "", "%v: could not parse claims", err)
		return ErrTokenNotValid
	}
	p.logf(id, "", "claims: %v", claims)

	if !claims.VerifyExpiresAt(now, true) {
		p.logerrorf(id, "", "%v", err)
		return ErrAuthFailed
	}
	if !claims.VerifyNotBefore(now, true) {
		p.logerrorf(id, "", "%v", err)
		return ErrAuthFailed
	}
	if claims["tid"] != id {
		p.logerrorf(id, "", "%v", err)
		return ErrAuthFailed
	}

	if claims["exp"].(float64)-claims["nbf"].(float64) > float64(monthUnix) {
		p.logerrorf(id, "", "jwt should not be valid for more than 31 days")
		return ErrAuthFailed
	}

	if claims.VerifyAudience(p.audience, false) {
		p.logerrorf(id, "", "%v", err)
		return ErrAuthFailed
	}

	return nil
}

// proxy logging utilities

// NOTE: cannot use logrus methods
func (p *proxy) logf(id string, remoteAddr string, format string, v ...interface{}) {
	p.logger.WithFields(logrus.Fields{
		"tunnel-id":   id,
		"remote-addr": remoteAddr,
	}).Printf(format, v...)
}

func (p *proxy) logerrorf(id string, remoteAddr string, format string, v ...interface{}) {
	p.logger.WithFields(logrus.Fields{
		"tunnel-id":   id,
		"remote-addr": remoteAddr,
	}).Errorf(format, v...)
}
<commit_msg>Fixed Audience verification<commit_after>package wsproxy

import (
	"bufio"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/gorilla/websocket"
	"github.com/taskcluster/websocktunnel/util"
	"github.com/taskcluster/websocktunnel/wsmux"

	"github.com/sirupsen/logrus"
	nullLog "github.com/sirupsen/logrus/hooks/test"
)

const (
	monthUnix = 31 * 24 * time.Hour
)

var (
	clientIdRe = regexp.MustCompile(`^[a-zA-Z0-9_~.%-]+$`)
)

// Config contains the run time parameters for the proxy
type Config struct {
	// Upgrader is a websocket.Upgrader instance which is used to upgrade incoming
	// websocket connections from Clients.
	Upgrader websocket.Upgrader

	// Logger is used to log proxy events. Refer util.Logger.
	Logger *logrus.Logger

	// JWTSecretA and JWTSecretB are used by the proxy to verify JWTs from Clients.
	JWTSecretA []byte
	JWTSecretB []byte

	// Domain and port where proxy will be hosted
	Domain string
	Port   int

	// set to true if serving with TLS
	TLS bool

	// Audience value for aud claim
	Audience string
}

// proxy is used to send http and ws requests to a registered client.
// New proxy can be created by using wsproxy.New()
type proxy struct {
	m               sync.RWMutex
	pool            map[string]*wsmux.Session
	upgrader        websocket.Upgrader
	logger          *logrus.Logger
	onSessionRemove func(string)
	jwtSecretA      []byte
	jwtSecretB      []byte
	domain          string
	port            int
	tls             bool
	audience        string
}

// New creates a new proxy instance and wraps it as an http.Handler.
func New(conf Config) (http.Handler, error) {
	return newProxy(conf)
}

// ServeHTTP implements http.Handler so that the proxy may be used as a handler in a Mux or http.Server
func (p *proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	p.logf("", r.RemoteAddr, "Host=%s Path=%s", r.Host, r.URL.Path)

	// Client registration requests are a GET of path / with some headers set
	if path, id := r.URL.Path, r.Header.Get("x-websocktunnel-id"); id != "" && path == "/" {
		tokenString := util.ExtractJWT(r.Header.Get("Authorization"))
		p.register(w, r, id, tokenString)
		return
	}

	// try to match viewer requests to https://<domain>/<id>/<path> (ignoring
	// the host field)
	s := strings.TrimPrefix(r.URL.RequestURI(), "/")
	index := strings.Index(s, "/")
	id, path := "", "/"
	if index < 0 {
		id = s
	} else {
		id = s[:index]
		path = s[index:]
	}
	if id == "" {
		http.NotFound(w, r)
		return
	}
	p.serveRequest(w, r, id, path)
}

func newProxy(conf Config) (*proxy, error) {
	p := &proxy{
		pool:       make(map[string]*wsmux.Session),
		upgrader:   conf.Upgrader,
		logger:     conf.Logger,
		jwtSecretA: conf.JWTSecretA,
		jwtSecretB: conf.JWTSecretB,
		domain:     conf.Domain,
		port:       conf.Port,
		tls:        conf.TLS,
		audience:   conf.Audience,
	}

	if conf.Port == 0 {
		panic("no port specified")
	}

	if len(p.jwtSecretA) == 0 || len(p.jwtSecretB) == 0 {
		panic("wsproxy: missing secrets")
	}

	if p.logger == nil {
		logger, _ := nullLog.NewNullLogger()
		p.logger = logger
	}

	return p, nil

}

// setSessionRemoveHandler sets a function which is called when a wsmux Session is removed from
// the proxy due to closure or error.
func (p *proxy) setSessionRemoveHandler(h func(string)) {
	p.m.Lock()
	defer p.m.Unlock()
	p.onSessionRemove = h
}

// getWorkerSession returns true if a session with the given id is present
func (p *proxy) getWorkerSession(id string) (*wsmux.Session, bool) {
	p.m.RLock()
	defer p.m.RUnlock()
	s, ok := p.pool[id]
	return s, ok
}

// removeTunnel is an idempotent operation which deletes a client session from the proxy's
// pool
func (p *proxy) removeTunnel(id string) {
	p.m.Lock()
	defer p.m.Unlock()
	delete(p.pool, id)
	p.logf(id, "", "session removed")
}

// register is used to connect a client to the proxy so that it can start serving API endpoints.
// The request must contain the tunnel ID in the url.
// The request is validated by the proxy and the http connection is upgraded to websocket.
func (p *proxy) register(w http.ResponseWriter, r *http.Request, id, tokenString string) {
	if !clientIdRe.MatchString(id) {
		p.logerrorf(id, r.RemoteAddr, "client ID is invalid")
		http.Error(w, http.StatusText(400), 400)
		return
	}

	p.logf(id, r.RemoteAddr, "requesting client registration")
	if tokenString == "" {
		// No jwt. Connection not authorized
		p.logerrorf(id, r.RemoteAddr, "could not retreive auth token")
		http.Error(w, http.StatusText(400), 400)
		return
	}

	if !websocket.IsWebSocketUpgrade(r) {
		p.logerrorf(id, r.RemoteAddr, "request must be websocket upgrade")
		http.NotFound(w, r)
		return
	}

	// validation does not require lock
	if err := p.validateJWT(id, tokenString); err != nil {
		p.logerrorf(id, r.RemoteAddr, "unable to validate token: %v", err)
		http.Error(w, http.StatusText(401), 401)
		return
	}

	p.m.Lock()

	// remove any existing session forcibly
	for {
		existingSession := p.pool[id]
		if existingSession == nil {
			break
		}

		// If there's an existing session, unlock, close the existing session
		// (which will remove it), and try again.  Note that Session.Close is
		// properly reentrant, so two goroutines calling this at the same time
		// will not cause an issue.
		p.m.Unlock()
		_ = existingSession.Close()
		p.m.Lock()
	}

	defer p.m.Unlock()

	delete(p.pool, id)

	header := make(http.Header)

	urlScheme := "http://"
	defaultPort := 80
	if p.tls {
		urlScheme = "https://"
		defaultPort = 443
	}
	var portStr string
	if p.port != defaultPort {
		portStr = fmt.Sprintf(":%d", p.port)
	}

	url := urlScheme + p.domain + portStr + "/" + id
	header.Set("x-websocktunnel-client-url", url)
	p.logf(id, r.RemoteAddr, "sending url= %s", url)
	conn, err := p.upgrader.Upgrade(w, r, header)
	if err != nil {
		p.logger.Print(err)
		return
	}

	// generate config
	conf := wsmux.Config{
		StreamBufferSize: 4 * 1024,
		CloseCallback: func() {
			p.removeTunnel(id)
			if p.onSessionRemove != nil {
				p.onSessionRemove(id)
			}
		},
		Log: p.logger,
	}

	p.pool[id] = wsmux.Server(conn, conf)
	p.logf(id, r.RemoteAddr, "added new tunnel")
}

// serveRequest serves tunnel endpoints to viewers
func (p *proxy) serveRequest(w http.ResponseWriter, r *http.Request, id string, path string) {
	// log new request arrival
	p.logf(id, r.RemoteAddr, "request: host=%s path=%s", r.Host, path)
	p.logf(id, r.RemoteAddr, "request: URL: %v", r.URL)

	session, ok := p.getWorkerSession(id)

	// return 504 (bad gateway) if tunnel is not registered on this proxy
	if !ok {
		p.logerrorf(id, r.RemoteAddr, "could not find requested tunnel")
		http.Error(w, fmt.Sprintf("No client is connected with that id"), 504)
		return
	}

	// set original path as header
	r.Header.Set("x-websocktunnel-original-path", r.URL.Path)

	// check for a websocket request
	if websocket.IsWebSocketUpgrade(r) {
		_ = p.websocketProxy(w, r, session, id)
		return
	}

	// Open a stream to the tunnel session
	// path is the modified RequestURI
	reqURI, err := url.ParseRequestURI(path)
	if err != nil {
		http.Error(w, http.StatusText(500), 500)
		return
	}
	r.URL = reqURI
	p.logf(id, r.RemoteAddr, "attempting to open new stream")
	reqStream, err := session.Open()
	if err != nil {
		p.logerrorf(id, r.RemoteAddr, "could not open stream: path=%s", path)
		http.Error(w, http.StatusText(500), 500)
		return
	}

	// rewrite path for tunnel and write request
	err = r.Write(reqStream)
	if err != nil {
		p.logerrorf(id, r.RemoteAddr, "could not write request: path=%s", path)
		http.Error(w, http.StatusText(500), 500)
		return
	}

	// read response from tunnel
	bufReader := bufio.NewReader(reqStream)
	resp, err := http.ReadResponse(bufReader, r)
	if err != nil {
		p.logerrorf(id, r.RemoteAddr, "could not read response: path=%s", path)
		http.Error(w, http.StatusText(500), 500)
		return
	}

	// manually proxy response
	// clear responseWriter headers and write response headers instead
	for k := range w.Header() {
		w.Header().Del(k)
	}
	for k, v := range resp.Header {
		w.Header()[k] = v
	}

	// dump headers
	w.WriteHeader(resp.StatusCode)
	if resp.Body == nil {
		return
	}

	// stream body to viewer and close wsmux stream
	defer func() {
		_ = reqStream.Close()
	}()

	flusher, ok := w.(http.Flusher)
	// flusher may not be implemented by a ResponseWriter wrapper
	// simple copy
	if !ok {
		n, err := io.Copy(w, resp.Body)
		p.logf(id, r.RemoteAddr, "data transfered over request: %d bytes, error: %v", n, err)
		// log here
		return
	}

	p.logf(id, r.RemoteAddr, "streaming http")
	wf := &threadSafeWriteFlusher{w: w, f: flusher}
	n, err := copyAndFlush(wf, resp.Body, 100*time.Millisecond)
	p.logf(id, r.RemoteAddr, "data transfered over request: %d bytes, error: %v", n, err)
}

// validate jwt
// jwt signing and verification algorithm must be HMAC
func (p *proxy) validateJWT(id string, tokenString string) error {
	// parse jwt token
	// default parser verifies iat token if present. This can be a problem because of clocks not being
	// in sync.
	parser := &jwt.Parser{
		ValidMethods:         []string{"HS256"},
		SkipClaimsValidation: true, // Claims will be verified if token can be decoded using secret
	}

	token, err := parser.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, ErrUnexpectedSigningMethod
		}
		return p.jwtSecretA, nil
	})

	if err != nil {
		// log first error
		p.logerrorf(id, "", "%v: trying with second secret", err)

		token, err = parser.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
			if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
				return nil, ErrUnexpectedSigningMethod
			}
			return p.jwtSecretB, nil
		})
	}

	if err != nil {
		p.logerrorf(id, "", "%v: auth failed", err)
		return ErrAuthFailed
	}

	// check claims
	now := time.Now().Unix()
	claims, ok := token.Claims.(jwt.MapClaims)

	if !ok {
		p.logerrorf(id, "", "%v: could not parse claims", err)
		return ErrTokenNotValid
	}
	p.logf(id, "", "claims: %v", claims)

	if !claims.VerifyExpiresAt(now, true) {
		p.logerrorf(id, "", "%v", err)
		return ErrAuthFailed
	}
	if !claims.VerifyNotBefore(now, true) {
		p.logerrorf(id, "", "%v", err)
		return ErrAuthFailed
	}
	if claims["tid"] != id {
		p.logerrorf(id, "", "%v", err)
		return ErrAuthFailed
	}

	if claims["exp"].(float64)-claims["nbf"].(float64) > float64(monthUnix) {
		p.logerrorf(id, "", "jwt should not be valid for more than 31 days")
		return ErrAuthFailed
	}

	if !claims.VerifyAudience(p.audience, false) {
		p.logerrorf(id, "", "%v", err)
		return ErrAuthFailed
	}

	return nil
}

// proxy logging utilities

// NOTE: cannot use logrus methods
func (p *proxy) logf(id string, remoteAddr string, format string, v ...interface{}) {
	p.logger.WithFields(logrus.Fields{
		"tunnel-id":   id,
		"remote-addr": remoteAddr,
	}).Printf(format, v...)
}

func (p *proxy) logerrorf(id string, remoteAddr string, format string, v ...interface{}) {
	p.logger.WithFields(logrus.Fields{
		"tunnel-id":   id,
		"remote-addr": remoteAddr,
	}).Errorf(format, v...)
}
<|endoftext|>


<commit_before>// This file was automatically generated. DO NOT EDIT.
// If you have any remark or suggestion do not hesitate to open an issue.

package account

import (
	"context"
	"reflect"

	"github.com/scaleway/scaleway-cli/internal/core"
	"github.com/scaleway/scaleway-sdk-go/api/account/v2alpha1"
	"github.com/scaleway/scaleway-sdk-go/scw"
)

// always import dependencies
var (
	_ = scw.RegionFrPar
)

func GetGeneratedCommands() *core.Commands {
	return core.NewCommands(
		accountRoot(),
		accountSSHKey(),
		accountSSHKeyList(),
		accountSSHKeyAdd(),
		accountSSHKeyGet(),
		accountSSHKeyUpdate(),
		accountSSHKeyRemove(),
	)
}
func accountRoot() *core.Command {
	return &core.Command{
		Short:     `This API allows to manage your scaleway account`,
		Long:      ``,
		Namespace: "account",
	}
}

func accountSSHKey() *core.Command {
	return &core.Command{
		Short:     `Manage your Scaleway SSH keys`,
		Long:      `Manage your Scaleway SSH keys.`,
		Namespace: "account",
		Resource:  "ssh-key",
	}
}

func accountSSHKeyList() *core.Command {
	return &core.Command{
		Short:     `List all SSH keys`,
		Long:      `List all SSH keys.`,
		Namespace: "account",
		Resource:  "ssh-key",
		Verb:      "list",
		ArgsType:  reflect.TypeOf(account.ListSSHKeysRequest{}),
		ArgSpecs: core.ArgSpecs{
			{
				Name:       "order-by",
				Required:   false,
				Positional: false,
				EnumValues: []string{"created_at_asc", "created_at_desc", "updated_at_asc", "updated_at_desc", "name_asc", "name_desc"},
			},
			{
				Name:       "name",
				Required:   false,
				Positional: false,
			},
			{
				Name:       "organization-id",
				Required:   false,
				Positional: false,
			},
		},
		Run: func(ctx context.Context, args interface{}) (i interface{}, e error) {
			request := args.(*account.ListSSHKeysRequest)

			client := core.ExtractClient(ctx)
			api := account.NewAPI(client)
			resp, err := api.ListSSHKeys(request, scw.WithAllPages())
			if err != nil {
				return nil, err
			}
			return resp.SSHKeys, nil

		},
		View: &core.View{Fields: []*core.ViewField{
			{
				FieldName: "id",
			},
			{
				FieldName: "name",
			},
			{
				FieldName: "created_at",
			},
			{
				FieldName: "updated_at",
			},
			{
				FieldName: "organization_id",
			},
			{
				FieldName: "creation_info.address",
			},
			{
				FieldName: "creation_info.country_code",
			},
			{
				FieldName: "creation_info.user_agent",
			},
		}},
	}
}

func accountSSHKeyAdd() *core.Command {
	return &core.Command{
		Short:     `Add a SSH key to your Scaleway account`,
		Long:      `Add a SSH key to your Scaleway account.`,
		Namespace: "account",
		Resource:  "ssh-key",
		Verb:      "add",
		ArgsType:  reflect.TypeOf(account.CreateSSHKeyRequest{}),
		ArgSpecs: core.ArgSpecs{
			{
				Name:       "name",
				Required:   false,
				Positional: false,
			},
			{
				Name:       "public-key",
				Required:   false,
				Positional: false,
			},
			core.OrganizationIDArgSpec(),
		},
		Run: func(ctx context.Context, args interface{}) (i interface{}, e error) {
			request := args.(*account.CreateSSHKeyRequest)

			client := core.ExtractClient(ctx)
			api := account.NewAPI(client)
			return api.CreateSSHKey(request)

		},
		Examples: []*core.Example{
			{
				Short:   "Add a given ssh key",
				Request: `{"name":"foobar","public_key":"ecdsa-sha2-nistp256 AAAAE2VjZHNhLXNoYTItbmlzdHAyNTYAAAAIbmlzdHAyNTYAAABBBBN+7q3Vx+9V6Ye8cfvV6wnXsQyp3wwy8ucSXYI3HJz9gB9x8lsDAQ8DUSz+sNajzKdhRtio1hNOyhn3O5ku5Kk= foobar@foobar"}`,
			},
		},
		SeeAlsos: []*core.SeeAlso{
			{
				Command: "scw account ssh-key list",
				Short:   "List all SSH keys",
			},
			{
				Command: "scw account ssh-key remove",
				Short:   "Remove an SSH key",
			},
		},
	}
}

func accountSSHKeyGet() *core.Command {
	return &core.Command{
		Short:     `Get SSH key details`,
		Long:      `Get SSH key details.`,
		Namespace: "account",
		Resource:  "ssh-key",
		Verb:      "get",
		ArgsType:  reflect.TypeOf(account.GetSSHKeyRequest{}),
		ArgSpecs: core.ArgSpecs{
			{
				Name:       "ssh-key-id",
				Required:   true,
				Positional: true,
			},
		},
		Run: func(ctx context.Context, args interface{}) (i interface{}, e error) {
			request := args.(*account.GetSSHKeyRequest)

			client := core.ExtractClient(ctx)
			api := account.NewAPI(client)
			return api.GetSSHKey(request)

		},
	}
}

func accountSSHKeyUpdate() *core.Command {
	return &core.Command{
		Short:     `Update an SSH key`,
		Long:      `Update an SSH key.`,
		Namespace: "account",
		Resource:  "ssh-key",
		Verb:      "update",
		ArgsType:  reflect.TypeOf(account.UpdateSSHKeyRequest{}),
		ArgSpecs: core.ArgSpecs{
			{
				Name:       "ssh-key-id",
				Required:   true,
				Positional: true,
			},
			{
				Name:       "name",
				Required:   false,
				Positional: false,
			},
		},
		Run: func(ctx context.Context, args interface{}) (i interface{}, e error) {
			request := args.(*account.UpdateSSHKeyRequest)

			client := core.ExtractClient(ctx)
			api := account.NewAPI(client)
			return api.UpdateSSHKey(request)

		},
	}
}

func accountSSHKeyRemove() *core.Command {
	return &core.Command{
		Short:     `Remove a SSH key from your Scaleway account`,
		Long:      `Remove a SSH key from your Scaleway account.`,
		Namespace: "account",
		Resource:  "ssh-key",
		Verb:      "remove",
		ArgsType:  reflect.TypeOf(account.DeleteSSHKeyRequest{}),
		ArgSpecs: core.ArgSpecs{
			{
				Name:       "ssh-key-id",
				Required:   true,
				Positional: true,
			},
		},
		Run: func(ctx context.Context, args interface{}) (i interface{}, e error) {
			request := args.(*account.DeleteSSHKeyRequest)

			client := core.ExtractClient(ctx)
			api := account.NewAPI(client)
			e = api.DeleteSSHKey(request)
			if e != nil {
				return nil, e
			}
			return &core.SuccessResult{
				Resource: "ssh-key",
				Verb:     "remove",
			}, nil
		},
		Examples: []*core.Example{
			{
				Short:   "Remove a given SSH key",
				Request: `{"ssh_key_id":"11111111-1111-1111-1111-111111111111"}`,
			},
		},
	}
}
<commit_msg>feat: update generated cli (#871)<commit_after>// This file was automatically generated. DO NOT EDIT.
// If you have any remark or suggestion do not hesitate to open an issue.

package account

import (
	"context"
	"reflect"

	"github.com/scaleway/scaleway-cli/internal/core"
	"github.com/scaleway/scaleway-sdk-go/api/account/v2alpha1"
	"github.com/scaleway/scaleway-sdk-go/scw"
)

// always import dependencies
var (
	_ = scw.RegionFrPar
)

func GetGeneratedCommands() *core.Commands {
	return core.NewCommands(
		accountRoot(),
		accountSSHKey(),
		accountSSHKeyList(),
		accountSSHKeyAdd(),
		accountSSHKeyGet(),
		accountSSHKeyUpdate(),
		accountSSHKeyRemove(),
	)
}
func accountRoot() *core.Command {
	return &core.Command{
		Short:     `This API allows to manage your scaleway account`,
		Long:      ``,
		Namespace: "account",
	}
}

func accountSSHKey() *core.Command {
	return &core.Command{
		Short:     `Manage your Scaleway SSH keys`,
		Long:      `Manage your Scaleway SSH keys.`,
		Namespace: "account",
		Resource:  "ssh-key",
	}
}

func accountSSHKeyList() *core.Command {
	return &core.Command{
		Short:     `List all SSH keys`,
		Long:      `List all SSH keys.`,
		Namespace: "account",
		Resource:  "ssh-key",
		Verb:      "list",
		ArgsType:  reflect.TypeOf(account.ListSSHKeysRequest{}),
		ArgSpecs: core.ArgSpecs{
			{
				Name:       "order-by",
				Required:   false,
				Positional: false,
				EnumValues: []string{"created_at_asc", "created_at_desc", "updated_at_asc", "updated_at_desc", "name_asc", "name_desc"},
			},
			{
				Name:       "name",
				Required:   false,
				Positional: false,
			},
			{
				Name:       "organization-id",
				Required:   false,
				Positional: false,
			},
		},
		Run: func(ctx context.Context, args interface{}) (i interface{}, e error) {
			request := args.(*account.ListSSHKeysRequest)

			client := core.ExtractClient(ctx)
			api := account.NewAPI(client)
			resp, err := api.ListSSHKeys(request, scw.WithAllPages())
			if err != nil {
				return nil, err
			}
			return resp.SSHKeys, nil

		},
		View: &core.View{Fields: []*core.ViewField{
			{
				FieldName: "id",
			},
			{
				FieldName: "name",
			},
			{
				FieldName: "created_at",
			},
			{
				FieldName: "updated_at",
			},
			{
				FieldName: "organization_id",
			},
			{
				FieldName: "creation_info.address",
			},
			{
				FieldName: "creation_info.country_code",
			},
			{
				FieldName: "creation_info.user_agent",
			},
		}},
	}
}

func accountSSHKeyAdd() *core.Command {
	return &core.Command{
		Short:     `Add a SSH key to your Scaleway account`,
		Long:      `Add a SSH key to your Scaleway account.`,
		Namespace: "account",
		Resource:  "ssh-key",
		Verb:      "add",
		ArgsType:  reflect.TypeOf(account.CreateSSHKeyRequest{}),
		ArgSpecs: core.ArgSpecs{
			{
				Name:       "name",
				Short:      `The name of the SSH key`,
				Required:   false,
				Positional: false,
			},
			{
				Name:       "public-key",
				Short:      `SSH public key. Currently ssh-rsa, ssh-dss (DSA), ssh-ed25519 and ecdsa keys with NIST curves are supported`,
				Required:   true,
				Positional: false,
			},
			core.OrganizationIDArgSpec(),
		},
		Run: func(ctx context.Context, args interface{}) (i interface{}, e error) {
			request := args.(*account.CreateSSHKeyRequest)

			client := core.ExtractClient(ctx)
			api := account.NewAPI(client)
			return api.CreateSSHKey(request)

		},
		Examples: []*core.Example{
			{
				Short:   "Add a given ssh key",
				Request: `null`,
			},
		},
		SeeAlsos: []*core.SeeAlso{
			{
				Command: "scw account ssh-key list",
				Short:   "List all SSH keys",
			},
			{
				Command: "scw account ssh-key remove",
				Short:   "Remove an SSH key",
			},
		},
	}
}

func accountSSHKeyGet() *core.Command {
	return &core.Command{
		Short:     `Get SSH key details`,
		Long:      `Get SSH key details.`,
		Namespace: "account",
		Resource:  "ssh-key",
		Verb:      "get",
		ArgsType:  reflect.TypeOf(account.GetSSHKeyRequest{}),
		ArgSpecs: core.ArgSpecs{
			{
				Name:       "ssh-key-id",
				Short:      `The ID of the SSH key`,
				Required:   true,
				Positional: true,
			},
		},
		Run: func(ctx context.Context, args interface{}) (i interface{}, e error) {
			request := args.(*account.GetSSHKeyRequest)

			client := core.ExtractClient(ctx)
			api := account.NewAPI(client)
			return api.GetSSHKey(request)

		},
	}
}

func accountSSHKeyUpdate() *core.Command {
	return &core.Command{
		Short:     `Update an SSH key`,
		Long:      `Update an SSH key.`,
		Namespace: "account",
		Resource:  "ssh-key",
		Verb:      "update",
		ArgsType:  reflect.TypeOf(account.UpdateSSHKeyRequest{}),
		ArgSpecs: core.ArgSpecs{
			{
				Name:       "ssh-key-id",
				Required:   true,
				Positional: true,
			},
			{
				Name:       "name",
				Required:   false,
				Positional: false,
			},
		},
		Run: func(ctx context.Context, args interface{}) (i interface{}, e error) {
			request := args.(*account.UpdateSSHKeyRequest)

			client := core.ExtractClient(ctx)
			api := account.NewAPI(client)
			return api.UpdateSSHKey(request)

		},
	}
}

func accountSSHKeyRemove() *core.Command {
	return &core.Command{
		Short:     `Remove a SSH key from your Scaleway account`,
		Long:      `Remove a SSH key from your Scaleway account.`,
		Namespace: "account",
		Resource:  "ssh-key",
		Verb:      "remove",
		ArgsType:  reflect.TypeOf(account.DeleteSSHKeyRequest{}),
		ArgSpecs: core.ArgSpecs{
			{
				Name:       "ssh-key-id",
				Required:   true,
				Positional: true,
			},
		},
		Run: func(ctx context.Context, args interface{}) (i interface{}, e error) {
			request := args.(*account.DeleteSSHKeyRequest)

			client := core.ExtractClient(ctx)
			api := account.NewAPI(client)
			e = api.DeleteSSHKey(request)
			if e != nil {
				return nil, e
			}
			return &core.SuccessResult{
				Resource: "ssh-key",
				Verb:     "remove",
			}, nil
		},
		Examples: []*core.Example{
			{
				Short:   "Remove a given SSH key",
				Request: `{"ssh_key_id":"11111111-1111-1111-1111-111111111111"}`,
			},
		},
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"fmt"
	"io"
	"os"
	"path"
	"time"

	"github.com/juju/errors"
	"github.com/juju/utils/filestorage"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"gopkg.in/mgo.v2/txn"

	"github.com/juju/juju/environs/storage"
	"github.com/juju/juju/state/backups/metadata"
	"github.com/juju/juju/version"
)

/*
Backups are not a part of juju state nor of normal state operations.
However, they certainly are tightly coupled with state (the very
subject of backups).  This puts backups in an odd position,
particularly with regard to the storage of backup metadata and
archives.  As a result, here are a couple concerns worth mentioning.

First, as noted above backup is about state but not a part of state.
So exposing backup-related methods on State would imply the wrong
thing.  Thus the backup functionality here in the state package (not
state/backups) is exposed as functions to which you pass a state
object.

Second, backup creates an archive file containing a dump of state's
mongo DB.  Storing backup metadata/archives in mongo is thus a
somewhat circular proposition that has the potential to cause
problems.  That may need further attention.

Note that state (and juju as a whole) currently does not have a
persistence layer abstraction to facilitate separating different
persistence needs and implementations.  As a consequence, state's
data, whether about how an environment should look or about existing
resources within an environment, is dumped essentially straight into
State's mongo connection.  The code in the state package does not
make any distinction between the two (nor does the package clearly
distinguish between state-related abstractions and state-related
data).

Backup adds yet another category, merely taking advantage of
State's DB.  In the interest of making the distinction clear, the
code that directly interacts with State (and its DB) lives in this
file.  As mentioned previously, the functionality here is exposed
through functions that take State, rather than as methods on State.
Furthermore, the bulk of the backup-related code, which does not need
direct interaction with State, lives in the state/backups package.
*/

// backupMetadataDoc is a mirror of metadata.Metadata, used just for DB storage.
type backupMetadataDoc struct {
	ID             string `bson:"_id"`
	Started        int64  `bson:"started,minsize"`
	Finished       int64  `bson:"finished,minsize"`
	Checksum       string `bson:"checksum"`
	ChecksumFormat string `bson:"checksumformat"`
	Size           int64  `bson:"size,minsize"`
	Stored         bool   `bson:"stored"`
	Notes          string `bson:"notes,omitempty"`

	// origin
	Environment string         `bson:"environment"`
	Machine     string         `bson:"machine"`
	Hostname    string         `bson:"hostname"`
	Version     version.Number `bson:"version"`
}

func (doc *backupMetadataDoc) fileSet() bool {
	if doc.Finished == 0 {
		return false
	}
	if doc.Checksum == "" {
		return false
	}
	if doc.ChecksumFormat == "" {
		return false
	}
	if doc.Size == 0 {
		return false
	}
	return true
}

func (doc *backupMetadataDoc) validate() error {
	if doc.ID == "" {
		return errors.New("missing ID")
	}
	if doc.Started == 0 {
		return errors.New("missing Started")
	}
	if doc.Environment == "" {
		return errors.New("missing Environment")
	}
	if doc.Machine == "" {
		return errors.New("missing Machine")
	}
	if doc.Hostname == "" {
		return errors.New("missing Hostname")
	}
	if doc.Version.Major == 0 {
		return errors.New("missing Version")
	}

	// Check the file-related fields.
	if !doc.fileSet() {
		if doc.Stored {
			return errors.New(`"Stored" flag is unexpectedly true`)
		}
		// Don't check the file-related fields.
		return nil
	}
	if doc.Finished == 0 {
		return errors.New("missing Finished")
	}
	if doc.Checksum == "" {
		return errors.New("missing Checksum")
	}
	if doc.ChecksumFormat == "" {
		return errors.New("missing ChecksumFormat")
	}
	if doc.Size == 0 {
		return errors.New("missing Size")
	}

	return nil
}

// asMetadata returns a new metadata.Metadata based on the backupMetadataDoc.
func (doc *backupMetadataDoc) asMetadata() *metadata.Metadata {
	// Create a new Metadata.
	origin := metadata.ExistingOrigin(
		doc.Environment,
		doc.Machine,
		doc.Hostname,
		doc.Version,
	)

	started := time.Unix(doc.Started, 0).UTC()
	meta := metadata.NewMetadata(
		*origin,
		doc.Notes,
		&started,
	)

	// The ID is already set.
	meta.SetID(doc.ID)

	// Exit early if file-related fields not set.
	if !doc.fileSet() {
		return meta
	}

	// Set the file-related fields.
	var finished *time.Time
	if doc.Finished != 0 {
		val := time.Unix(doc.Finished, 0).UTC()
		finished = &val
	}
	err := meta.Finish(doc.Size, doc.Checksum, doc.ChecksumFormat, finished)
	if err != nil {
		// The doc should have already been validated.  An error here
		// indicates that Metadata changed and backupMetadataDoc did not
		// accommodate the change.  Thus an error here indicates a
		// developer "error".  A caller should not need to worry about
		// that case so we panic instead of passing the error out.
		panic(fmt.Sprintf("unexpectedly invalid metadata doc: %v", err))
	}
	if doc.Stored {
		meta.SetStored()
	}
	return meta
}

// updateFromMetadata copies the corresponding data from the backup
// Metadata into the backupMetadataDoc.
func (doc *backupMetadataDoc) updateFromMetadata(metadata *metadata.Metadata) {
	finished := metadata.Finished()
	// Ignore metadata.ID.
	doc.Started = metadata.Started().Unix()
	if finished != nil {
		doc.Finished = finished.Unix()
	}
	doc.Checksum = metadata.Checksum()
	doc.ChecksumFormat = metadata.ChecksumFormat()
	doc.Size = metadata.Size()
	doc.Stored = metadata.Stored()
	doc.Notes = metadata.Notes()

	origin := metadata.Origin()
	doc.Environment = origin.Environment()
	doc.Machine = origin.Machine()
	doc.Hostname = origin.Hostname()
	doc.Version = origin.Version()
}

//---------------------------
// DB operations

// getBackupMetadata returns the backup metadata associated with "id".
// If "id" does not match any stored records, an error satisfying
// juju/errors.IsNotFound() is returned.
func getBackupMetadata(st *State, id string) (*metadata.Metadata, error) {
	collection, closer := st.getCollection(backupsMetaC)
	defer closer()

	var doc backupMetadataDoc
	// There can only be one!
	err := collection.FindId(id).One(&doc)
	if err == mgo.ErrNotFound {
		return nil, errors.NotFoundf("backup metadata %q", id)
	} else if err != nil {
		return nil, errors.Annotate(err, "error getting backup metadata")
	}

	if err := doc.validate(); err != nil {
		return nil, errors.Trace(err)
	}
	return doc.asMetadata(), nil
}

// newBackupID returns a new ID for a state backup.  The format is the
// UTC timestamp from the metadata followed by the environment ID:
// "YYYYMMDD-hhmmss.<env ID>".  This makes the ID a little more human-
// consumable (in contrast to a plain UUID string).  Ideally we would
// use some form of environment name rather than the UUID, but for now
// the raw env ID is sufficient.
func newBackupID(metadata *metadata.Metadata) string {
	rawts := metadata.Started()
	Y, M, D := rawts.Date()
	h, m, s := rawts.Clock()
	timestamp := fmt.Sprintf("%04d%02d%02d-%02d%02d%02d", Y, M, D, h, m, s)
	origin := metadata.Origin()
	env := origin.Environment()
	return timestamp + "." + env
}

// addBackupMetadata stores metadata for a backup where it can be
// accessed later.  It returns a new ID that is associated with the
// backup.  If the provided metadata already has an ID set, it is
// ignored.
func addBackupMetadata(st *State, metadata *metadata.Metadata) (string, error) {
	// We use our own mongo _id value since the auto-generated one from
	// mongo may contain sensitive data (see bson.ObjectID).
	id := newBackupID(metadata)
	return id, addBackupMetadataID(st, metadata, id)
}

func addBackupMetadataID(st *State, metadata *metadata.Metadata, id string) error {
	var doc backupMetadataDoc
	doc.updateFromMetadata(metadata)
	doc.ID = id
	if err := doc.validate(); err != nil {
		return errors.Trace(err)
	}

	ops := []txn.Op{{
		C:      backupsMetaC,
		Id:     doc.ID,
		Assert: txn.DocMissing,
		Insert: doc,
	}}
	if err := st.runTransaction(ops); err != nil {
		if err == txn.ErrAborted {
			return errors.AlreadyExistsf("backup metadata %q", doc.ID)
		}
		return errors.Annotate(err, "error running transaction")
	}

	return nil
}

// setBackupStored updates the backup metadata associated with "id"
// to indicate that a backup archive has been stored.  If "id" does
// not match any stored records, an error satisfying
// juju/errors.IsNotFound() is returned.
func setBackupStored(st *State, id string) error {
	ops := []txn.Op{{
		C:      backupsMetaC,
		Id:     id,
		Assert: txn.DocExists,
		Update: bson.D{{"$set", bson.D{
			{"stored", true},
		}}},
	}}
	if err := st.runTransaction(ops); err != nil {
		if err == txn.ErrAborted {
			return errors.NotFoundf(id)
		}
		return errors.Annotate(err, "error running transaction")
	}
	return nil
}

//---------------------------
// metadata storage

// NewBackupsOrigin returns a snapshot of where backup was run.  That
// snapshot is a new backup Origin value, for use in a backup's
// metadata.  Every value except for the machine name is populated
// either from juju state or some other implicit mechanism.
func NewBackupsOrigin(st *State, machine string) *metadata.Origin {
	// hostname could be derived from the environment...
	hostname, err := os.Hostname()
	if err != nil {
		// If os.Hostname() is not working, something is woefully wrong.
		// Run for the hills.
		panic(fmt.Sprintf("could not get hostname (system unstable?): %v", err))
	}
	origin := metadata.NewOrigin(
		st.EnvironTag().Id(),
		machine,
		hostname,
	)
	return origin
}

// Ensure we satisfy the interface.
var _ = filestorage.MetadataStorage((*backupMetadataStorage)(nil))

type backupMetadataStorage struct {
	state *State
}

func newBackupMetadataStorage(st *State) filestorage.MetadataStorage {
	stor := backupMetadataStorage{
		state: st,
	}
	return &stor
}

func (s *backupMetadataStorage) AddDoc(doc interface{}) (string, error) {
	meta, ok := doc.(*metadata.Metadata)
	if !ok {
		return "", errors.Errorf("doc must be of type state.backups.metadata.Metadata")
	}
	return addBackupMetadata(s.state, meta)
}

func (s *backupMetadataStorage) Doc(id string) (interface{}, error) {
	return s.Metadata(id)
}

func (s *backupMetadataStorage) Metadata(id string) (filestorage.Metadata, error) {
	metadata, err := getBackupMetadata(s.state, id)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return metadata, nil
}

func (s *backupMetadataStorage) ListDocs() ([]interface{}, error) {
	metas, err := s.ListMetadata()
	if err != nil {
		return nil, errors.Trace(err)
	}
	docs := []interface{}{}
	for _, meta := range metas {
		docs = append(docs, meta)
	}
	return docs, nil
}

func (s *backupMetadataStorage) ListMetadata() ([]filestorage.Metadata, error) {
	// This will be implemented when backups needs this functionality.
	// For now the method is stubbed out for the same of the
	// MetadataStorage interface.
	return nil, errors.NotImplementedf("ListMetadata")
}

func (s *backupMetadataStorage) RemoveDoc(id string) error {
	// This will be implemented when backups needs this functionality.
	// For now the method is stubbed out for the same of the
	// MetadataStorage interface.
	return errors.NotImplementedf("RemoveDoc")
}

func (s *backupMetadataStorage) New() filestorage.Metadata {
	origin := NewBackupsOrigin(s.state, "")
	return metadata.NewMetadata(*origin, "", nil)
}

func (s *backupMetadataStorage) SetStored(meta filestorage.Metadata) error {
	err := setBackupStored(s.state, meta.ID())
	if err != nil {
		return errors.Trace(err)
	}
	meta.SetStored()
	return nil
}

//---------------------------
// raw file storage

const backupStorageRoot = ""

// Ensure we satisfy the interface.
var _ filestorage.RawFileStorage = (*envFileStorage)(nil)

type envFileStorage struct {
	envStor storage.Storage
	root    string
}

func newBackupFileStorage(envStor storage.Storage, root string) filestorage.RawFileStorage {
	// Due to circular imports we cannot simply get the storage from
	// State using environs.GetStorage().
	stor := envFileStorage{
		envStor: envStor,
		root:    root,
	}
	return &stor
}

func (s *envFileStorage) path(id string) string {
	// Use of path.Join instead of filepath.Join is intentional - this
	// is an environment storage path not a filesystem path.
	return path.Join(s.root, id)
}

func (s *envFileStorage) File(id string) (io.ReadCloser, error) {
	return s.envStor.Get(s.path(id))
}

func (s *envFileStorage) AddFile(id string, file io.Reader, size int64) error {
	return s.envStor.Put(s.path(id), file, size)
}

func (s *envFileStorage) RemoveFile(id string) error {
	return s.envStor.Remove(s.path(id))
}

//---------------------------
// backup storage

// NewBackupsStorage returns a new FileStorage to use for storing backup
// archives (and metadata).
func NewBackupsStorage(st *State, envStor storage.Storage) filestorage.FileStorage {
	files := newBackupFileStorage(envStor, backupStorageRoot)
	docs := newBackupMetadataStorage(st)
	return filestorage.NewFileStorage(docs, files)
}
<commit_msg>Use a proper storage root.<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"fmt"
	"io"
	"os"
	"path"
	"time"

	"github.com/juju/errors"
	"github.com/juju/utils/filestorage"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"gopkg.in/mgo.v2/txn"

	"github.com/juju/juju/environs/storage"
	"github.com/juju/juju/state/backups/metadata"
	"github.com/juju/juju/version"
)

/*
Backups are not a part of juju state nor of normal state operations.
However, they certainly are tightly coupled with state (the very
subject of backups).  This puts backups in an odd position,
particularly with regard to the storage of backup metadata and
archives.  As a result, here are a couple concerns worth mentioning.

First, as noted above backup is about state but not a part of state.
So exposing backup-related methods on State would imply the wrong
thing.  Thus the backup functionality here in the state package (not
state/backups) is exposed as functions to which you pass a state
object.

Second, backup creates an archive file containing a dump of state's
mongo DB.  Storing backup metadata/archives in mongo is thus a
somewhat circular proposition that has the potential to cause
problems.  That may need further attention.

Note that state (and juju as a whole) currently does not have a
persistence layer abstraction to facilitate separating different
persistence needs and implementations.  As a consequence, state's
data, whether about how an environment should look or about existing
resources within an environment, is dumped essentially straight into
State's mongo connection.  The code in the state package does not
make any distinction between the two (nor does the package clearly
distinguish between state-related abstractions and state-related
data).

Backup adds yet another category, merely taking advantage of
State's DB.  In the interest of making the distinction clear, the
code that directly interacts with State (and its DB) lives in this
file.  As mentioned previously, the functionality here is exposed
through functions that take State, rather than as methods on State.
Furthermore, the bulk of the backup-related code, which does not need
direct interaction with State, lives in the state/backups package.
*/

// backupMetadataDoc is a mirror of metadata.Metadata, used just for DB storage.
type backupMetadataDoc struct {
	ID             string `bson:"_id"`
	Started        int64  `bson:"started,minsize"`
	Finished       int64  `bson:"finished,minsize"`
	Checksum       string `bson:"checksum"`
	ChecksumFormat string `bson:"checksumformat"`
	Size           int64  `bson:"size,minsize"`
	Stored         bool   `bson:"stored"`
	Notes          string `bson:"notes,omitempty"`

	// origin
	Environment string         `bson:"environment"`
	Machine     string         `bson:"machine"`
	Hostname    string         `bson:"hostname"`
	Version     version.Number `bson:"version"`
}

func (doc *backupMetadataDoc) fileSet() bool {
	if doc.Finished == 0 {
		return false
	}
	if doc.Checksum == "" {
		return false
	}
	if doc.ChecksumFormat == "" {
		return false
	}
	if doc.Size == 0 {
		return false
	}
	return true
}

func (doc *backupMetadataDoc) validate() error {
	if doc.ID == "" {
		return errors.New("missing ID")
	}
	if doc.Started == 0 {
		return errors.New("missing Started")
	}
	if doc.Environment == "" {
		return errors.New("missing Environment")
	}
	if doc.Machine == "" {
		return errors.New("missing Machine")
	}
	if doc.Hostname == "" {
		return errors.New("missing Hostname")
	}
	if doc.Version.Major == 0 {
		return errors.New("missing Version")
	}

	// Check the file-related fields.
	if !doc.fileSet() {
		if doc.Stored {
			return errors.New(`"Stored" flag is unexpectedly true`)
		}
		// Don't check the file-related fields.
		return nil
	}
	if doc.Finished == 0 {
		return errors.New("missing Finished")
	}
	if doc.Checksum == "" {
		return errors.New("missing Checksum")
	}
	if doc.ChecksumFormat == "" {
		return errors.New("missing ChecksumFormat")
	}
	if doc.Size == 0 {
		return errors.New("missing Size")
	}

	return nil
}

// asMetadata returns a new metadata.Metadata based on the backupMetadataDoc.
func (doc *backupMetadataDoc) asMetadata() *metadata.Metadata {
	// Create a new Metadata.
	origin := metadata.ExistingOrigin(
		doc.Environment,
		doc.Machine,
		doc.Hostname,
		doc.Version,
	)

	started := time.Unix(doc.Started, 0).UTC()
	meta := metadata.NewMetadata(
		*origin,
		doc.Notes,
		&started,
	)

	// The ID is already set.
	meta.SetID(doc.ID)

	// Exit early if file-related fields not set.
	if !doc.fileSet() {
		return meta
	}

	// Set the file-related fields.
	var finished *time.Time
	if doc.Finished != 0 {
		val := time.Unix(doc.Finished, 0).UTC()
		finished = &val
	}
	err := meta.Finish(doc.Size, doc.Checksum, doc.ChecksumFormat, finished)
	if err != nil {
		// The doc should have already been validated.  An error here
		// indicates that Metadata changed and backupMetadataDoc did not
		// accommodate the change.  Thus an error here indicates a
		// developer "error".  A caller should not need to worry about
		// that case so we panic instead of passing the error out.
		panic(fmt.Sprintf("unexpectedly invalid metadata doc: %v", err))
	}
	if doc.Stored {
		meta.SetStored()
	}
	return meta
}

// updateFromMetadata copies the corresponding data from the backup
// Metadata into the backupMetadataDoc.
func (doc *backupMetadataDoc) updateFromMetadata(metadata *metadata.Metadata) {
	finished := metadata.Finished()
	// Ignore metadata.ID.
	doc.Started = metadata.Started().Unix()
	if finished != nil {
		doc.Finished = finished.Unix()
	}
	doc.Checksum = metadata.Checksum()
	doc.ChecksumFormat = metadata.ChecksumFormat()
	doc.Size = metadata.Size()
	doc.Stored = metadata.Stored()
	doc.Notes = metadata.Notes()

	origin := metadata.Origin()
	doc.Environment = origin.Environment()
	doc.Machine = origin.Machine()
	doc.Hostname = origin.Hostname()
	doc.Version = origin.Version()
}

//---------------------------
// DB operations

// getBackupMetadata returns the backup metadata associated with "id".
// If "id" does not match any stored records, an error satisfying
// juju/errors.IsNotFound() is returned.
func getBackupMetadata(st *State, id string) (*metadata.Metadata, error) {
	collection, closer := st.getCollection(backupsMetaC)
	defer closer()

	var doc backupMetadataDoc
	// There can only be one!
	err := collection.FindId(id).One(&doc)
	if err == mgo.ErrNotFound {
		return nil, errors.NotFoundf("backup metadata %q", id)
	} else if err != nil {
		return nil, errors.Annotate(err, "error getting backup metadata")
	}

	if err := doc.validate(); err != nil {
		return nil, errors.Trace(err)
	}
	return doc.asMetadata(), nil
}

// newBackupID returns a new ID for a state backup.  The format is the
// UTC timestamp from the metadata followed by the environment ID:
// "YYYYMMDD-hhmmss.<env ID>".  This makes the ID a little more human-
// consumable (in contrast to a plain UUID string).  Ideally we would
// use some form of environment name rather than the UUID, but for now
// the raw env ID is sufficient.
func newBackupID(metadata *metadata.Metadata) string {
	rawts := metadata.Started()
	Y, M, D := rawts.Date()
	h, m, s := rawts.Clock()
	timestamp := fmt.Sprintf("%04d%02d%02d-%02d%02d%02d", Y, M, D, h, m, s)
	origin := metadata.Origin()
	env := origin.Environment()
	return timestamp + "." + env
}

// addBackupMetadata stores metadata for a backup where it can be
// accessed later.  It returns a new ID that is associated with the
// backup.  If the provided metadata already has an ID set, it is
// ignored.
func addBackupMetadata(st *State, metadata *metadata.Metadata) (string, error) {
	// We use our own mongo _id value since the auto-generated one from
	// mongo may contain sensitive data (see bson.ObjectID).
	id := newBackupID(metadata)
	return id, addBackupMetadataID(st, metadata, id)
}

func addBackupMetadataID(st *State, metadata *metadata.Metadata, id string) error {
	var doc backupMetadataDoc
	doc.updateFromMetadata(metadata)
	doc.ID = id
	if err := doc.validate(); err != nil {
		return errors.Trace(err)
	}

	ops := []txn.Op{{
		C:      backupsMetaC,
		Id:     doc.ID,
		Assert: txn.DocMissing,
		Insert: doc,
	}}
	if err := st.runTransaction(ops); err != nil {
		if err == txn.ErrAborted {
			return errors.AlreadyExistsf("backup metadata %q", doc.ID)
		}
		return errors.Annotate(err, "error running transaction")
	}

	return nil
}

// setBackupStored updates the backup metadata associated with "id"
// to indicate that a backup archive has been stored.  If "id" does
// not match any stored records, an error satisfying
// juju/errors.IsNotFound() is returned.
func setBackupStored(st *State, id string) error {
	ops := []txn.Op{{
		C:      backupsMetaC,
		Id:     id,
		Assert: txn.DocExists,
		Update: bson.D{{"$set", bson.D{
			{"stored", true},
		}}},
	}}
	if err := st.runTransaction(ops); err != nil {
		if err == txn.ErrAborted {
			return errors.NotFoundf(id)
		}
		return errors.Annotate(err, "error running transaction")
	}
	return nil
}

//---------------------------
// metadata storage

// NewBackupsOrigin returns a snapshot of where backup was run.  That
// snapshot is a new backup Origin value, for use in a backup's
// metadata.  Every value except for the machine name is populated
// either from juju state or some other implicit mechanism.
func NewBackupsOrigin(st *State, machine string) *metadata.Origin {
	// hostname could be derived from the environment...
	hostname, err := os.Hostname()
	if err != nil {
		// If os.Hostname() is not working, something is woefully wrong.
		// Run for the hills.
		panic(fmt.Sprintf("could not get hostname (system unstable?): %v", err))
	}
	origin := metadata.NewOrigin(
		st.EnvironTag().Id(),
		machine,
		hostname,
	)
	return origin
}

// Ensure we satisfy the interface.
var _ = filestorage.MetadataStorage((*backupMetadataStorage)(nil))

type backupMetadataStorage struct {
	state *State
}

func newBackupMetadataStorage(st *State) filestorage.MetadataStorage {
	stor := backupMetadataStorage{
		state: st,
	}
	return &stor
}

func (s *backupMetadataStorage) AddDoc(doc interface{}) (string, error) {
	meta, ok := doc.(*metadata.Metadata)
	if !ok {
		return "", errors.Errorf("doc must be of type state.backups.metadata.Metadata")
	}
	return addBackupMetadata(s.state, meta)
}

func (s *backupMetadataStorage) Doc(id string) (interface{}, error) {
	return s.Metadata(id)
}

func (s *backupMetadataStorage) Metadata(id string) (filestorage.Metadata, error) {
	metadata, err := getBackupMetadata(s.state, id)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return metadata, nil
}

func (s *backupMetadataStorage) ListDocs() ([]interface{}, error) {
	metas, err := s.ListMetadata()
	if err != nil {
		return nil, errors.Trace(err)
	}
	docs := []interface{}{}
	for _, meta := range metas {
		docs = append(docs, meta)
	}
	return docs, nil
}

func (s *backupMetadataStorage) ListMetadata() ([]filestorage.Metadata, error) {
	// This will be implemented when backups needs this functionality.
	// For now the method is stubbed out for the same of the
	// MetadataStorage interface.
	return nil, errors.NotImplementedf("ListMetadata")
}

func (s *backupMetadataStorage) RemoveDoc(id string) error {
	// This will be implemented when backups needs this functionality.
	// For now the method is stubbed out for the same of the
	// MetadataStorage interface.
	return errors.NotImplementedf("RemoveDoc")
}

func (s *backupMetadataStorage) New() filestorage.Metadata {
	origin := NewBackupsOrigin(s.state, "")
	return metadata.NewMetadata(*origin, "", nil)
}

func (s *backupMetadataStorage) SetStored(meta filestorage.Metadata) error {
	err := setBackupStored(s.state, meta.ID())
	if err != nil {
		return errors.Trace(err)
	}
	meta.SetStored()
	return nil
}

//---------------------------
// raw file storage

const backupStorageRoot = "backups"

// Ensure we satisfy the interface.
var _ filestorage.RawFileStorage = (*envFileStorage)(nil)

type envFileStorage struct {
	envStor storage.Storage
	root    string
}

func newBackupFileStorage(envStor storage.Storage, root string) filestorage.RawFileStorage {
	// Due to circular imports we cannot simply get the storage from
	// State using environs.GetStorage().
	stor := envFileStorage{
		envStor: envStor,
		root:    root,
	}
	return &stor
}

func (s *envFileStorage) path(id string) string {
	// Use of path.Join instead of filepath.Join is intentional - this
	// is an environment storage path not a filesystem path.
	return path.Join(s.root, id)
}

func (s *envFileStorage) File(id string) (io.ReadCloser, error) {
	return s.envStor.Get(s.path(id))
}

func (s *envFileStorage) AddFile(id string, file io.Reader, size int64) error {
	return s.envStor.Put(s.path(id), file, size)
}

func (s *envFileStorage) RemoveFile(id string) error {
	return s.envStor.Remove(s.path(id))
}

//---------------------------
// backup storage

// NewBackupsStorage returns a new FileStorage to use for storing backup
// archives (and metadata).
func NewBackupsStorage(st *State, envStor storage.Storage) filestorage.FileStorage {
	files := newBackupFileStorage(envStor, backupStorageRoot)
	docs := newBackupMetadataStorage(st)
	return filestorage.NewFileStorage(docs, files)
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2015 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"bytes"
	"fmt"
	"io"

	"github.com/FactomProject/FactomCode/util"
	"github.com/davecgh/go-spew/spew"
)

// defaultTransactionAlloc is the default size used for the backing array
// for transactions.  The transaction array will dynamically grow as needed, but
// this figure is intended to provide enough space for the number of
// transactions in the vast majority of blocks without needing to grow the
// backing array multiple times.
const defaultTransactionAlloc = 2048

// MaxBlocksPerMsg is the maximum number of blocks allowed per message.
const MaxBlocksPerMsg = 10

// MaxBlockPayload is the maximum bytes a block message can be in bytes.
const MaxBlockPayload = 1000000 // Not actually 1MB which would be 1024 * 1024

// maxTxPerBlock is the maximum number of transactions that could
// possibly fit into a block.
const maxTxPerBlock = (MaxBlockPayload / minTxPayload) + 1

// TxLoc holds locator data for the offset and length of where a transaction is
// located within a MsgBlock data buffer.
type TxLoc struct {
	TxStart int
	TxLen   int
}

// MsgBlock implements the Message interface and represents a bitcoin
// block message.  It is used to deliver block and transaction information in
// response to a getdata message (MsgGetData) for a given block hash.
type MsgBlock struct {
	Header       BlockHeader
	Transactions []*MsgTx
}

// AddTransaction adds a transaction to the message.
func (msg *MsgBlock) AddTransaction(tx *MsgTx) error {
	msg.Transactions = append(msg.Transactions, tx)
	return nil

}

// ClearTransactions removes all transactions from the message.
func (msg *MsgBlock) ClearTransactions() {
	msg.Transactions = make([]*MsgTx, 0, defaultTransactionAlloc)
}

// BtcDecode decodes r using the bitcoin protocol encoding into the receiver.
// This is part of the Message interface implementation.
// See Deserialize for decoding blocks stored to disk, such as in a database, as
// opposed to decoding blocks from the wire.
func (msg *MsgBlock) BtcDecode(r io.Reader, pver uint32) error {
	err := readBlockHeader(r, pver, &msg.Header)
	if err != nil {
		return err
	}

	txCount, err := readVarInt(r, pver)
	if err != nil {
		return err
	}

	// Prevent more transactions than could possibly fit into a block.
	// It would be possible to cause memory exhaustion and panics without
	// a sane upper bound on this count.
	if txCount > maxTxPerBlock {
		str := fmt.Sprintf("too many transactions to fit into a block "+
			"[count %d, max %d]", txCount, maxTxPerBlock)
		return messageError("MsgBlock.BtcDecode", str)
	}

	msg.Transactions = make([]*MsgTx, 0, txCount)
	for i := uint64(0); i < txCount; i++ {
		tx := MsgTx{}
		err := tx.BtcDecode(r, pver)
		if err != nil {
			return err
		}
		msg.Transactions = append(msg.Transactions, &tx)
	}

	return nil
}

// Deserialize decodes a block from r into the receiver using a format that is
// suitable for long-term storage such as a database while respecting the
// Version field in the block.  This function differs from BtcDecode in that
// BtcDecode decodes from the bitcoin wire protocol as it was sent across the
// network.  The wire encoding can technically differ depending on the protocol
// version and doesn't even really need to match the format of a stored block at
// all.  As of the time this comment was written, the encoded block is the same
// in both instances, but there is a distinct difference and separating the two
// allows the API to be flexible enough to deal with changes.
func (msg *MsgBlock) Deserialize(r io.Reader) error {
	// At the current time, there is no difference between the wire encoding
	// at protocol version 0 and the stable long-term storage format.  As
	// a result, make use of BtcDecode.
	return msg.BtcDecode(r, 0)
}

// DeserializeTxLoc decodes r in the same manner Deserialize does, but it takes
// a byte buffer instead of a generic reader and returns a slice containing the start and length of
// each transaction within the raw data that is being deserialized.
func (msg *MsgBlock) DeserializeTxLoc(r *bytes.Buffer) ([]TxLoc, error) {
	util.Trace()
	fullLen := r.Len()

	fmt.Println("fullLen=", fullLen, spew.Sdump(r.Bytes()))

	// At the current time, there is no difference between the wire encoding
	// at protocol version 0 and the stable long-term storage format.  As
	// a result, make use of existing wire protocol functions.
	err := readBlockHeader(r, 0, &msg.Header)
	if err != nil {
		return nil, err
	}
	util.Trace()

	txCount, err := readVarInt(r, 0)
	if err != nil {
		return nil, err
	}

	util.Trace(fmt.Sprintf("txCount= %d", txCount))

	// Prevent more transactions than could possibly fit into a block.
	// It would be possible to cause memory exhaustion and panics without
	// a sane upper bound on this count.
	if txCount > maxTxPerBlock {
		str := fmt.Sprintf("too many transactions to fit into a block "+
			"[count %d, max %d]", txCount, maxTxPerBlock)
		return nil, messageError("MsgBlock.DeserializeTxLoc", str)
	}
	util.Trace()

	// Deserialize each transaction while keeping track of its location
	// within the byte stream.
	msg.Transactions = make([]*MsgTx, 0, txCount)
	txLocs := make([]TxLoc, txCount)
	for i := uint64(0); i < txCount; i++ {
		txLocs[i].TxStart = fullLen - r.Len()
		tx := MsgTx{}
		err := tx.Deserialize(r)
		if err != nil {
			return nil, err
		}

		util.Trace("tx deserialized= " + spew.Sdump(tx))

		msg.Transactions = append(msg.Transactions, &tx)
		txLocs[i].TxLen = (fullLen - r.Len()) - txLocs[i].TxStart
	}

	return txLocs, nil
}

// BtcEncode encodes the receiver to w using the bitcoin protocol encoding.
// This is part of the Message interface implementation.
// See Serialize for encoding blocks to be stored to disk, such as in a
// database, as opposed to encoding blocks for the wire.
func (msg *MsgBlock) BtcEncode(w io.Writer, pver uint32) error {
	err := writeBlockHeader(w, pver, &msg.Header)
	if err != nil {
		return err
	}

	err = writeVarInt(w, pver, uint64(len(msg.Transactions)))
	if err != nil {
		return err
	}

	for _, tx := range msg.Transactions {
		err = tx.BtcEncode(w, pver)
		if err != nil {
			return err
		}
	}

	return nil
}

// Serialize encodes the block to w using a format that suitable for long-term
// storage such as a database while respecting the Version field in the block.
// This function differs from BtcEncode in that BtcEncode encodes the block to
// the bitcoin wire protocol in order to be sent across the network.  The wire
// encoding can technically differ depending on the protocol version and doesn't
// even really need to match the format of a stored block at all.  As of the
// time this comment was written, the encoded block is the same in both
// instances, but there is a distinct difference and separating the two allows
// the API to be flexible enough to deal with changes.
func (msg *MsgBlock) Serialize(w io.Writer) error {
	// At the current time, there is no difference between the wire encoding
	// at protocol version 0 and the stable long-term storage format.  As
	// a result, make use of BtcEncode.
	return msg.BtcEncode(w, 0)
}

// SerializeSize returns the number of bytes it would take to serialize the
// the block.
func (msg *MsgBlock) SerializeSize() int {
	// Block header bytes + Serialized varint size for the number of
	// transactions.
	n := blockHeaderLen + VarIntSerializeSize(uint64(len(msg.Transactions)))

	for _, tx := range msg.Transactions {
		n += tx.SerializeSize()
	}

	return n
}

// Command returns the protocol command string for the message.  This is part
// of the Message interface implementation.
func (msg *MsgBlock) Command() string {
	return CmdBlock
}

// MaxPayloadLength returns the maximum length the payload can be for the
// receiver.  This is part of the Message interface implementation.
func (msg *MsgBlock) MaxPayloadLength(pver uint32) uint32 {
	// Block header at 80 bytes + transaction count + max transactions
	// which can vary up to the MaxBlockPayload (including the block header
	// and transaction count).
	return MaxBlockPayload
}

// BlockSha computes the block identifier hash for this block.
func (msg *MsgBlock) BlockSha() (ShaHash, error) {
	return msg.Header.BlockSha()
}

// TxShas returns a slice of hashes of all of transactions in this block.
func (msg *MsgBlock) TxShas() ([]ShaHash, error) {
	shaList := make([]ShaHash, 0, len(msg.Transactions))
	for _, tx := range msg.Transactions {
		// Ignore error here since TxSha can't fail in the current
		// implementation except due to run-time panics.
		sha, _ := tx.TxSha()
		shaList = append(shaList, sha)
	}
	return shaList, nil
}

// NewMsgBlock returns a new bitcoin block message that conforms to the
// Message interface.  See MsgBlock for details.
func NewMsgBlock(blockHeader *BlockHeader) *MsgBlock {
	return &MsgBlock{
		Header:       *blockHeader,
		Transactions: make([]*MsgTx, 0, defaultTransactionAlloc),
	}
}
<commit_msg>Changed the batch size back to 500<commit_after>// Copyright (c) 2013-2015 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"bytes"
	"fmt"
	"io"

	"github.com/FactomProject/FactomCode/util"
	"github.com/davecgh/go-spew/spew"
)

// defaultTransactionAlloc is the default size used for the backing array
// for transactions.  The transaction array will dynamically grow as needed, but
// this figure is intended to provide enough space for the number of
// transactions in the vast majority of blocks without needing to grow the
// backing array multiple times.
const defaultTransactionAlloc = 2048

// MaxBlocksPerMsg is the maximum number of blocks allowed per message.
const MaxBlocksPerMsg = 500

// MaxBlockPayload is the maximum bytes a block message can be in bytes.
const MaxBlockPayload = 1000000 // Not actually 1MB which would be 1024 * 1024

// maxTxPerBlock is the maximum number of transactions that could
// possibly fit into a block.
const maxTxPerBlock = (MaxBlockPayload / minTxPayload) + 1

// TxLoc holds locator data for the offset and length of where a transaction is
// located within a MsgBlock data buffer.
type TxLoc struct {
	TxStart int
	TxLen   int
}

// MsgBlock implements the Message interface and represents a bitcoin
// block message.  It is used to deliver block and transaction information in
// response to a getdata message (MsgGetData) for a given block hash.
type MsgBlock struct {
	Header       BlockHeader
	Transactions []*MsgTx
}

// AddTransaction adds a transaction to the message.
func (msg *MsgBlock) AddTransaction(tx *MsgTx) error {
	msg.Transactions = append(msg.Transactions, tx)
	return nil

}

// ClearTransactions removes all transactions from the message.
func (msg *MsgBlock) ClearTransactions() {
	msg.Transactions = make([]*MsgTx, 0, defaultTransactionAlloc)
}

// BtcDecode decodes r using the bitcoin protocol encoding into the receiver.
// This is part of the Message interface implementation.
// See Deserialize for decoding blocks stored to disk, such as in a database, as
// opposed to decoding blocks from the wire.
func (msg *MsgBlock) BtcDecode(r io.Reader, pver uint32) error {
	err := readBlockHeader(r, pver, &msg.Header)
	if err != nil {
		return err
	}

	txCount, err := readVarInt(r, pver)
	if err != nil {
		return err
	}

	// Prevent more transactions than could possibly fit into a block.
	// It would be possible to cause memory exhaustion and panics without
	// a sane upper bound on this count.
	if txCount > maxTxPerBlock {
		str := fmt.Sprintf("too many transactions to fit into a block "+
			"[count %d, max %d]", txCount, maxTxPerBlock)
		return messageError("MsgBlock.BtcDecode", str)
	}

	msg.Transactions = make([]*MsgTx, 0, txCount)
	for i := uint64(0); i < txCount; i++ {
		tx := MsgTx{}
		err := tx.BtcDecode(r, pver)
		if err != nil {
			return err
		}
		msg.Transactions = append(msg.Transactions, &tx)
	}

	return nil
}

// Deserialize decodes a block from r into the receiver using a format that is
// suitable for long-term storage such as a database while respecting the
// Version field in the block.  This function differs from BtcDecode in that
// BtcDecode decodes from the bitcoin wire protocol as it was sent across the
// network.  The wire encoding can technically differ depending on the protocol
// version and doesn't even really need to match the format of a stored block at
// all.  As of the time this comment was written, the encoded block is the same
// in both instances, but there is a distinct difference and separating the two
// allows the API to be flexible enough to deal with changes.
func (msg *MsgBlock) Deserialize(r io.Reader) error {
	// At the current time, there is no difference between the wire encoding
	// at protocol version 0 and the stable long-term storage format.  As
	// a result, make use of BtcDecode.
	return msg.BtcDecode(r, 0)
}

// DeserializeTxLoc decodes r in the same manner Deserialize does, but it takes
// a byte buffer instead of a generic reader and returns a slice containing the start and length of
// each transaction within the raw data that is being deserialized.
func (msg *MsgBlock) DeserializeTxLoc(r *bytes.Buffer) ([]TxLoc, error) {
	util.Trace()
	fullLen := r.Len()

	fmt.Println("fullLen=", fullLen, spew.Sdump(r.Bytes()))

	// At the current time, there is no difference between the wire encoding
	// at protocol version 0 and the stable long-term storage format.  As
	// a result, make use of existing wire protocol functions.
	err := readBlockHeader(r, 0, &msg.Header)
	if err != nil {
		return nil, err
	}
	util.Trace()

	txCount, err := readVarInt(r, 0)
	if err != nil {
		return nil, err
	}

	util.Trace(fmt.Sprintf("txCount= %d", txCount))

	// Prevent more transactions than could possibly fit into a block.
	// It would be possible to cause memory exhaustion and panics without
	// a sane upper bound on this count.
	if txCount > maxTxPerBlock {
		str := fmt.Sprintf("too many transactions to fit into a block "+
			"[count %d, max %d]", txCount, maxTxPerBlock)
		return nil, messageError("MsgBlock.DeserializeTxLoc", str)
	}
	util.Trace()

	// Deserialize each transaction while keeping track of its location
	// within the byte stream.
	msg.Transactions = make([]*MsgTx, 0, txCount)
	txLocs := make([]TxLoc, txCount)
	for i := uint64(0); i < txCount; i++ {
		txLocs[i].TxStart = fullLen - r.Len()
		tx := MsgTx{}
		err := tx.Deserialize(r)
		if err != nil {
			return nil, err
		}

		util.Trace("tx deserialized= " + spew.Sdump(tx))

		msg.Transactions = append(msg.Transactions, &tx)
		txLocs[i].TxLen = (fullLen - r.Len()) - txLocs[i].TxStart
	}

	return txLocs, nil
}

// BtcEncode encodes the receiver to w using the bitcoin protocol encoding.
// This is part of the Message interface implementation.
// See Serialize for encoding blocks to be stored to disk, such as in a
// database, as opposed to encoding blocks for the wire.
func (msg *MsgBlock) BtcEncode(w io.Writer, pver uint32) error {
	err := writeBlockHeader(w, pver, &msg.Header)
	if err != nil {
		return err
	}

	err = writeVarInt(w, pver, uint64(len(msg.Transactions)))
	if err != nil {
		return err
	}

	for _, tx := range msg.Transactions {
		err = tx.BtcEncode(w, pver)
		if err != nil {
			return err
		}
	}

	return nil
}

// Serialize encodes the block to w using a format that suitable for long-term
// storage such as a database while respecting the Version field in the block.
// This function differs from BtcEncode in that BtcEncode encodes the block to
// the bitcoin wire protocol in order to be sent across the network.  The wire
// encoding can technically differ depending on the protocol version and doesn't
// even really need to match the format of a stored block at all.  As of the
// time this comment was written, the encoded block is the same in both
// instances, but there is a distinct difference and separating the two allows
// the API to be flexible enough to deal with changes.
func (msg *MsgBlock) Serialize(w io.Writer) error {
	// At the current time, there is no difference between the wire encoding
	// at protocol version 0 and the stable long-term storage format.  As
	// a result, make use of BtcEncode.
	return msg.BtcEncode(w, 0)
}

// SerializeSize returns the number of bytes it would take to serialize the
// the block.
func (msg *MsgBlock) SerializeSize() int {
	// Block header bytes + Serialized varint size for the number of
	// transactions.
	n := blockHeaderLen + VarIntSerializeSize(uint64(len(msg.Transactions)))

	for _, tx := range msg.Transactions {
		n += tx.SerializeSize()
	}

	return n
}

// Command returns the protocol command string for the message.  This is part
// of the Message interface implementation.
func (msg *MsgBlock) Command() string {
	return CmdBlock
}

// MaxPayloadLength returns the maximum length the payload can be for the
// receiver.  This is part of the Message interface implementation.
func (msg *MsgBlock) MaxPayloadLength(pver uint32) uint32 {
	// Block header at 80 bytes + transaction count + max transactions
	// which can vary up to the MaxBlockPayload (including the block header
	// and transaction count).
	return MaxBlockPayload
}

// BlockSha computes the block identifier hash for this block.
func (msg *MsgBlock) BlockSha() (ShaHash, error) {
	return msg.Header.BlockSha()
}

// TxShas returns a slice of hashes of all of transactions in this block.
func (msg *MsgBlock) TxShas() ([]ShaHash, error) {
	shaList := make([]ShaHash, 0, len(msg.Transactions))
	for _, tx := range msg.Transactions {
		// Ignore error here since TxSha can't fail in the current
		// implementation except due to run-time panics.
		sha, _ := tx.TxSha()
		shaList = append(shaList, sha)
	}
	return shaList, nil
}

// NewMsgBlock returns a new bitcoin block message that conforms to the
// Message interface.  See MsgBlock for details.
func NewMsgBlock(blockHeader *BlockHeader) *MsgBlock {
	return &MsgBlock{
		Header:       *blockHeader,
		Transactions: make([]*MsgTx, 0, defaultTransactionAlloc),
	}
}
<|endoftext|>