<commit_before>// Copyright 2016 The Gosl Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package io

import (
	"bufio"
	"bytes"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/cpmech/gosl/chk"
)

// FnKey returns the file name key (without path and extension, if any)
func FnKey(fn string) string {
	base := filepath.Base(fn)
	return base[:len(base)-len(filepath.Ext(base))]
}

// FnExt returns the extension of a file name.
// The extension is the suffix beginning at the final dot in the final element of path; it is empty if there is no dot.
func FnExt(fn string) string {
	return filepath.Ext(fn)
}

// PathKey returs the full path except the extension
func PathKey(fn string) string {
	return fn[:len(fn)-len(filepath.Ext(fn))]
}

// AppendToFile appends data to an existent (or new) file
func AppendToFile(fn string, buffer ...*bytes.Buffer) {
	fil, err := os.OpenFile(os.ExpandEnv(fn), os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		chk.Panic(_fileio_err01, fn)
	}
	defer fil.Close()
	for k, _ := range buffer {
		if buffer[k] != nil {
			fil.Write(buffer[k].Bytes())
		}
	}
}

// WriteFile writes data to a new file
func WriteFile(fn string, buffer ...*bytes.Buffer) {
	//fil, err := os.OpenFile(os.ExpandEnv(fn), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
	fil, err := os.Create(os.ExpandEnv(fn))
	if err != nil {
		chk.Panic(_fileio_err02, fn)
	}
	defer fil.Close()
	for k, _ := range buffer {
		if buffer[k] != nil {
			fil.Write(buffer[k].Bytes())
		}
	}
}

// WriteFileD writes data to a new file and create a directory to save the file within
func WriteFileD(dirout, fn string, buffer ...*bytes.Buffer) {
	os.MkdirAll(dirout, 0777)
	WriteFile(filepath.Join(dirout, fn), buffer...)
}

// WriteFileS writes data string to a new file
func WriteFileS(fn, data string) {
	fil, err := os.Create(os.ExpandEnv(fn))
	if err != nil {
		chk.Panic(_fileio_err02, fn)
	}
	defer fil.Close()
	fil.WriteString(data)
}

// WriteFileSD writes data string to a new file after creating a directory to save the file
func WriteFileSD(dirout, fn, data string) {
	os.MkdirAll(dirout, 0777)
	WriteFileS(filepath.Join(dirout, fn), data)
}

// WriteFileV writes data to a new file (and shows message: file written)
func WriteFileV(fn string, buffer ...*bytes.Buffer) {
	WriteFile(fn, buffer...)
	Pf("File <")
	PfBlue(fn)
	Pf("> written\n")
}

// WriteFileVD writes data to a new file (and shows message: file written), after creating a directory
func WriteFileVD(dirout, fn string, buffer ...*bytes.Buffer) {
	WriteFileD(dirout, fn, buffer...)
	Pf("file <%s> written\n", filepath.Join(dirout, fn))
}

// WriteBytesToFile writes bytes to a new file
func WriteBytesToFile(fn string, b []byte) {
	fil, err := os.Create(os.ExpandEnv(fn))
	if err != nil {
		chk.Panic(_fileio_err03, fn)
	}
	defer fil.Close()
	if _, err = fil.Write(b); err != nil {
		chk.Panic(_fileio_err04, err.Error())
	}
}

// WriteBytesToFile writes bytes to a new file and create a directory to save the file within
func WriteBytesToFileD(dirout, fn string, b []byte) {
	os.MkdirAll(dirout, 0777)
	WriteBytesToFile(filepath.Join(dirout, fn), b)
}

// OpenFileR opens a file for reading data
func OpenFileR(fn string) (fil *os.File, err error) {
	fil, err = os.Open(os.ExpandEnv(fn))
	return
}

// ReadFile reads bytes from a file
func ReadFile(fn string) (b []byte, err error) {
	return ioutil.ReadFile(os.ExpandEnv(fn))
}

// ReadLinesCallback is used in ReadLines to process line by line during reading of a file
type ReadLinesCallback func(idx int, line string) (stop bool)

// ReadLines reads lines from a file and calls ReadLinesCallback to process each line being read
func ReadLines(fn string, cb ReadLinesCallback) {
	fil, err := os.Open(os.ExpandEnv(fn))
	if err != nil {
		chk.Panic(_fileio_err05, fn)
	}
	defer fil.Close()
	r := bufio.NewReader(fil)
	idx := 0
	for {
		lin, prefix, errl := r.ReadLine()
		if prefix {
			chk.Panic(_fileio_err06, fn)
		}
		if errl == io.EOF {
			break
		}
		if errl != nil {
			chk.Panic(_fileio_err07, fn)
		}
		stop := cb(idx, string(lin))
		if stop {
			break
		}
		idx++
	}
}

// ReadLinesFile reads lines from a file and calls ReadLinesCallback to process each line being read
func ReadLinesFile(fil *os.File, cb ReadLinesCallback) (oserr error) {
	r := bufio.NewReader(fil)
	idx := 0
	for {
		lin, prefix, errl := r.ReadLine()
		if prefix {
			return chk.Err(_fileio_err08, fil.Name())
		}
		if errl == io.EOF {
			break
		}
		if errl != nil {
			return chk.Err(_fileio_err09, fil.Name())
		}
		stop := cb(idx, string(lin))
		if stop {
			break
		}
		idx++
	}
	return
}

// RemoveAll deletes all files matching filename specified by key (be careful)
func RemoveAll(key string) {
	fns, _ := filepath.Glob(os.ExpandEnv(key))
	for _, fn := range fns {
		os.RemoveAll(fn)
	}
}

//  ReadTableOrPanic reads text file as ReadTable; but panic on errors
func ReadTableOrPanic(fn string) (keys []string, T map[string][]float64) {
	var err error
	keys, T, err = ReadTable(fn)
	if err != nil {
		chk.Panic("cannot read table:\n%v", err)
	}
	return
}

// ReadTable reads a text file in which the first line contains the headers and the next lines the float64
// type of numeric values. The number of columns must be equal, including for the headers
func ReadTable(fn string) (keys []string, T map[string][]float64, err error) {
	f, err := OpenFileR(fn)
	if err != nil {
		return
	}
	header := true
	ReadLinesFile(f, func(idx int, line string) (stop bool) {
		r := strings.Fields(line)
		if len(r) == 0 { // skip empty lines
			return
		}
		if r[0] == "#" { // skip comments
			return
		}
		ncol := len(r)
		if ncol < 1 {
			err = chk.Err("ReadTable: number of columns must be at least 1")
			return true
		}
		if header {
			T = make(map[string][]float64)
			keys = make([]string, ncol)
			for i := 0; i < ncol; i++ {
				keys[i] = r[i]
				T[r[i]] = make([]float64, 0)
			}
			header = false
		} else {
			for i := 0; i < ncol; i++ {
				T[keys[i]] = append(T[keys[i]], Atof(r[i]))
			}
		}
		return
	})
	return
}

// ReadMatrix reads a text file in which the float64 type of numeric values represent
// a matrix of data. The number of columns must be equal, including for the headers
func ReadMatrix(fn string) (M [][]float64, err error) {
	f, err := OpenFileR(fn)
	if err != nil {
		return
	}
	ncol_fix := 0
	ReadLinesFile(f, func(idx int, line string) (stop bool) {
		r := strings.Fields(line)
		if len(r) == 0 { // skip empty lines
			return
		}
		if r[0] == "#" { // skip comments
			return
		}
		ncol := len(r)
		if ncol < 1 {
			err = chk.Err("ReadMatrix: number of columns must be at least 1")
			return true
		}
		if M == nil {
			M = make([][]float64, 0)
			ncol_fix = ncol
		}
		if ncol != ncol_fix {
			err = chk.Err("ReadMatrix: nubmer of columns must be equal for all lines (%d != %d)", ncol, ncol_fix)
			return
		}
		vals := make([]float64, ncol)
		for i := 0; i < ncol; i++ {
			vals[i] = Atof(r[i])
		}
		M = append(M, vals)
		return
	})
	return
}

// error messages
var (
	_fileio_err01 = "cannot create file <%s>"
	_fileio_err02 = "cannot create file <%s>"
	_fileio_err03 = "cannot create file <%s>"
	_fileio_err04 = "failed with error: %s"
	_fileio_err05 = "could not open file <%s>"
	_fileio_err06 = "cannot read long line yet. file = <%s>"
	_fileio_err07 = "cannot read line. file = <%s>"
	_fileio_err08 = "cannot read long line yet. file = <%s>"
	_fileio_err09 = "cannot read line. file = <%s>"
)
<commit_msg>Remove variables with error messages in fileio<commit_after>// Copyright 2016 The Gosl Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package io

import (
	"bufio"
	"bytes"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/cpmech/gosl/chk"
)

// FnKey returns the file name key (without path and extension, if any)
func FnKey(fn string) string {
	base := filepath.Base(fn)
	return base[:len(base)-len(filepath.Ext(base))]
}

// FnExt returns the extension of a file name.
// The extension is the suffix beginning at the final dot in the final element of path; it is empty if there is no dot.
func FnExt(fn string) string {
	return filepath.Ext(fn)
}

// PathKey returs the full path except the extension
func PathKey(fn string) string {
	return fn[:len(fn)-len(filepath.Ext(fn))]
}

// AppendToFile appends data to an existent (or new) file
func AppendToFile(fn string, buffer ...*bytes.Buffer) {
	fil, err := os.OpenFile(os.ExpandEnv(fn), os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		chk.Panic("cannot create file <%s>", fn)
	}
	defer fil.Close()
	for k, _ := range buffer {
		if buffer[k] != nil {
			fil.Write(buffer[k].Bytes())
		}
	}
}

// WriteFile writes data to a new file
func WriteFile(fn string, buffer ...*bytes.Buffer) {
	fil, err := os.Create(os.ExpandEnv(fn))
	if err != nil {
		chk.Panic("cannot create file <%s>", fn)
	}
	defer fil.Close()
	for k, _ := range buffer {
		if buffer[k] != nil {
			fil.Write(buffer[k].Bytes())
		}
	}
}

// WriteFileD writes data to a new file and create a directory to save the file within
func WriteFileD(dirout, fn string, buffer ...*bytes.Buffer) {
	os.MkdirAll(dirout, 0777)
	WriteFile(filepath.Join(dirout, fn), buffer...)
}

// WriteFileS writes data string to a new file
func WriteFileS(fn, data string) {
	fil, err := os.Create(os.ExpandEnv(fn))
	if err != nil {
		chk.Panic("cannot create file <%s>", fn)
	}
	defer fil.Close()
	fil.WriteString(data)
}

// WriteFileSD writes data string to a new file after creating a directory to save the file
func WriteFileSD(dirout, fn, data string) {
	os.MkdirAll(dirout, 0777)
	WriteFileS(filepath.Join(dirout, fn), data)
}

// WriteFileV writes data to a new file (and shows message: file written)
func WriteFileV(fn string, buffer ...*bytes.Buffer) {
	WriteFile(fn, buffer...)
	Pf("file <%s> written\n", fn)
}

// WriteFileVD writes data to a new file (and shows message: file written), after creating a directory
func WriteFileVD(dirout, fn string, buffer ...*bytes.Buffer) {
	WriteFileD(dirout, fn, buffer...)
	Pf("file <%s> written\n", filepath.Join(dirout, fn))
}

// WriteBytesToFile writes bytes to a new file
func WriteBytesToFile(fn string, b []byte) {
	fil, err := os.Create(os.ExpandEnv(fn))
	if err != nil {
		chk.Panic("cannot create file <%s>", fn)
	}
	defer fil.Close()
	if _, err = fil.Write(b); err != nil {
		chk.Panic("%v", err)
	}
}

// WriteBytesToFile writes bytes to a new file and create a directory to save the file within
func WriteBytesToFileD(dirout, fn string, b []byte) {
	os.MkdirAll(dirout, 0777)
	WriteBytesToFile(filepath.Join(dirout, fn), b)
}

// OpenFileR opens a file for reading data
func OpenFileR(fn string) (fil *os.File, err error) {
	fil, err = os.Open(os.ExpandEnv(fn))
	return
}

// ReadFile reads bytes from a file
func ReadFile(fn string) (b []byte, err error) {
	return ioutil.ReadFile(os.ExpandEnv(fn))
}

// ReadLinesCallback is used in ReadLines to process line by line during reading of a file
type ReadLinesCallback func(idx int, line string) (stop bool)

// ReadLines reads lines from a file and calls ReadLinesCallback to process each line being read
func ReadLines(fn string, cb ReadLinesCallback) {
	fil, err := os.Open(os.ExpandEnv(fn))
	if err != nil {
		chk.Panic("could not open file <%s>", fn)
	}
	defer fil.Close()
	r := bufio.NewReader(fil)
	idx := 0
	for {
		lin, prefix, errl := r.ReadLine()
		if prefix {
			chk.Panic("cannot read long line. file = <%s>", fn)
		}
		if errl == io.EOF {
			break
		}
		if errl != nil {
			chk.Panic("cannot read line. file = <%s>", fn)
		}
		stop := cb(idx, string(lin))
		if stop {
			break
		}
		idx++
	}
}

// ReadLinesFile reads lines from a file and calls ReadLinesCallback to process each line being read
func ReadLinesFile(fil *os.File, cb ReadLinesCallback) (oserr error) {
	r := bufio.NewReader(fil)
	idx := 0
	for {
		lin, prefix, errl := r.ReadLine()
		if prefix {
			return chk.Err("cannot read long line. file = <%s>", fil.Name())
		}
		if errl == io.EOF {
			break
		}
		if errl != nil {
			return chk.Err("cannot read line. file = <%s>", fil.Name())
		}
		stop := cb(idx, string(lin))
		if stop {
			break
		}
		idx++
	}
	return
}

// RemoveAll deletes all files matching filename specified by key (be careful)
func RemoveAll(key string) {
	fns, _ := filepath.Glob(os.ExpandEnv(key))
	for _, fn := range fns {
		os.RemoveAll(fn)
	}
}

//  ReadTableOrPanic reads text file as ReadTable; but panic on errors
func ReadTableOrPanic(fn string) (keys []string, T map[string][]float64) {
	var err error
	keys, T, err = ReadTable(fn)
	if err != nil {
		chk.Panic("cannot read table:\n%v", err)
	}
	return
}

// ReadTable reads a text file in which the first line contains the headers and the next lines the float64
// type of numeric values. The number of columns must be equal, including for the headers
func ReadTable(fn string) (keys []string, T map[string][]float64, err error) {
	f, err := OpenFileR(fn)
	if err != nil {
		return
	}
	header := true
	ReadLinesFile(f, func(idx int, line string) (stop bool) {
		r := strings.Fields(line)
		if len(r) == 0 { // skip empty lines
			return
		}
		if r[0] == "#" { // skip comments
			return
		}
		ncol := len(r)
		if ncol < 1 {
			err = chk.Err("ReadTable: number of columns must be at least 1")
			return true
		}
		if header {
			T = make(map[string][]float64)
			keys = make([]string, ncol)
			for i := 0; i < ncol; i++ {
				keys[i] = r[i]
				T[r[i]] = make([]float64, 0)
			}
			header = false
		} else {
			for i := 0; i < ncol; i++ {
				T[keys[i]] = append(T[keys[i]], Atof(r[i]))
			}
		}
		return
	})
	return
}

// ReadMatrix reads a text file in which the float64 type of numeric values represent
// a matrix of data. The number of columns must be equal, including for the headers
func ReadMatrix(fn string) (M [][]float64, err error) {
	f, err := OpenFileR(fn)
	if err != nil {
		return
	}
	ncol_fix := 0
	ReadLinesFile(f, func(idx int, line string) (stop bool) {
		r := strings.Fields(line)
		if len(r) == 0 { // skip empty lines
			return
		}
		if r[0] == "#" { // skip comments
			return
		}
		ncol := len(r)
		if ncol < 1 {
			err = chk.Err("ReadMatrix: number of columns must be at least 1")
			return true
		}
		if M == nil {
			M = make([][]float64, 0)
			ncol_fix = ncol
		}
		if ncol != ncol_fix {
			err = chk.Err("ReadMatrix: nubmer of columns must be equal for all lines (%d != %d)", ncol, ncol_fix)
			return
		}
		vals := make([]float64, ncol)
		for i := 0; i < ncol; i++ {
			vals[i] = Atof(r[i])
		}
		M = append(M, vals)
		return
	})
	return
}
<|endoftext|>


<commit_before>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * io/reader.go                                           *
 *                                                        *
 * hprose reader for Go.                                  *
 *                                                        *
 * LastModified: Sep 14, 2016                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package io

import (
	"bytes"
	"errors"
	"math/big"
	"reflect"
	"time"

	"github.com/hprose/hprose-golang/util"
)

// Reader is a fine-grained operation struct for Hprose unserialization
// when JSONCompatible is true, the Map data will unserialize to map[string]interface as the default type
type Reader struct {
	RawReader
	Simple         bool
	fieldsRef      [][]*fieldCache
	ref            []interface{}
	JSONCompatible bool
}

// NewReader is the constructor for Hprose Reader
func NewReader(buf []byte, simple bool) (reader *Reader) {
	reader = new(Reader)
	reader.buf = buf
	reader.Simple = simple
	return
}

// Unserialize a data from the reader
func (r *Reader) Unserialize(p interface{}) {
	v := reflect.ValueOf(p)
	if v.Kind() != reflect.Ptr {
		panic(errors.New("Unserialize: argument p must be a pointer"))
	}
	r.ReadValue(v.Elem())
}

// ReadValue from the reader
func (r *Reader) ReadValue(v reflect.Value) {
	tag := r.readByte()
	decoder := valueDecoders[v.Kind()]
	decoder(r, v, tag)
}

// CheckTag the next byte in reader is the expected tag or not
func (r *Reader) CheckTag(expectTag byte) (tag byte) {
	tag = r.readByte()
	if tag != expectTag {
		unexpectedTag(tag, []byte{expectTag})
	}
	return
}

// CheckTags the next byte in reader in the expected tags
func (r *Reader) CheckTags(expectTags []byte) (tag byte) {
	tag = r.readByte()
	if bytes.IndexByte(expectTags, tag) == -1 {
		unexpectedTag(tag, expectTags)
	}
	return
}

// ReadBool from the reader
func (r *Reader) ReadBool() bool {
	tag := r.readByte()
	decoder := boolDecoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "bool")
	return false
}

// ReadIntWithoutTag from the reader
func (r *Reader) ReadIntWithoutTag() int {
	return readInt(&r.ByteReader)
}

// ReadInt from the reader
func (r *Reader) ReadInt() int64 {
	tag := r.readByte()
	decoder := intDecoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "int64")
	return 0
}

// ReadUint from the reader
func (r *Reader) ReadUint() uint64 {
	tag := r.readByte()
	decoder := uintDecoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "uint64")
	return 0
}

// ReadFloat32 from the reader
func (r *Reader) ReadFloat32() float32 {
	tag := r.readByte()
	decoder := float32Decoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "float32")
	return 0
}

// ReadFloat64 from the reader
func (r *Reader) ReadFloat64() float64 {
	tag := r.readByte()
	decoder := float64Decoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "float64")
	return 0
}

// ReadComplex64 from the reader
func (r *Reader) ReadComplex64() complex64 {
	tag := r.readByte()
	decoder := complex64Decoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "complex64")
	return 0
}

// ReadComplex128 from the reader
func (r *Reader) ReadComplex128() complex128 {
	tag := r.readByte()
	decoder := complex128Decoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "complex128")
	return 0
}

// ReadStringWithoutTag from the reader
func (r *Reader) ReadStringWithoutTag() (str string) {
	str = readString(&r.ByteReader)
	if !r.Simple {
		setReaderRef(r, str)
	}
	return str
}

// ReadString from the reader
func (r *Reader) ReadString() (str string) {
	tag := r.readByte()
	decoder := stringDecoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "string")
	return ""
}

// ReadBytesWithoutTag from the reader
func (r *Reader) ReadBytesWithoutTag() (b []byte) {
	l := readLength(&r.ByteReader)
	b = make([]byte, l)
	if _, err := r.Read(b); err != nil {
		panic(err)
	}
	r.readByte()
	if !r.Simple {
		setReaderRef(r, b)
	}
	return
}

// ReadDateTimeWithoutTag from the reader
func (r *Reader) ReadDateTimeWithoutTag() (dt time.Time) {
	year := r.read4Digit()
	month := r.read2Digit()
	day := r.read2Digit()
	tag := r.readByte()
	var hour, min, sec, nsec int
	if tag == TagTime {
		hour = r.read2Digit()
		min = r.read2Digit()
		sec = r.read2Digit()
		tag = r.readByte()
		if tag == TagPoint {
			nsec, tag = readNsec(&r.ByteReader)
		}
	}
	var loc *time.Location
	if tag == TagUTC {
		loc = time.UTC
	} else {
		loc = time.Local
	}
	dt = time.Date(year, time.Month(month), day, hour, min, sec, nsec, loc)
	if !r.Simple {
		setReaderRef(r, &dt)
	}
	return
}

// ReadTimeWithoutTag from the reader
func (r *Reader) ReadTimeWithoutTag() (t time.Time) {
	hour := r.read2Digit()
	min := r.read2Digit()
	sec := r.read2Digit()
	tag := r.readByte()
	var nsec int
	if tag == TagPoint {
		nsec, tag = readNsec(&r.ByteReader)
	}
	var loc *time.Location
	if tag == TagUTC {
		loc = time.UTC
	} else {
		loc = time.Local
	}
	t = time.Date(1970, 1, 1, hour, min, sec, nsec, loc)
	if !r.Simple {
		setReaderRef(r, &t)
	}
	return
}

// ReadBigIntWithoutTag from the reader
func (r *Reader) ReadBigIntWithoutTag() *big.Int {
	b := readUntil(&r.ByteReader, TagSemicolon)
	i, _ := new(big.Int).SetString(util.ByteString(b), 10)
	return i
}

// ReadSliceWithoutTag from the reader
func (r *Reader) ReadSliceWithoutTag() []reflect.Value {
	l := r.ReadCount()
	v := make([]reflect.Value, l)
	if !r.Simple {
		setReaderRef(r, nil)
	}
	for i := 0; i < l; i++ {
		v[i] = reflect.New(interfaceType).Elem()
		r.ReadValue(v[i])
	}
	r.readByte()
	return v
}

// ReadSlice from the reader
func (r *Reader) ReadSlice(v []reflect.Value) {
	l := len(v)
	if !r.Simple {
		setReaderRef(r, nil)
	}
	for i := 0; i < l; i++ {
		r.ReadValue(v[i])
	}
	r.readByte()
}

// ReadCount of array, slice, map or struct field
func (r *Reader) ReadCount() int {
	return int(ReadInt64(&r.ByteReader, TagOpenbrace))
}

// Reset the reference counter
func (r *Reader) Reset() {
	if r.fieldsRef != nil {
		r.fieldsRef = r.fieldsRef[:0]
	}
	if r.Simple {
		return
	}
	if r.ref != nil {
		r.ref = r.ref[:0]
	}
}

// private methods & functions

func (r *Reader) readRef() interface{} {
	if r.Simple {
		panic(errors.New("reference unserialization can't support in simple mode"))
	}
	return readRef(r, readInt(&r.ByteReader))
}

func setReaderRef(r *Reader, o interface{}) {
	r.ref = append(r.ref, o)
}

func readRef(r *Reader, i int) interface{} {
	return r.ref[i]
}

var tagStringMap = map[byte]string{
	'0':         "int",
	'1':         "int",
	'2':         "int",
	'3':         "int",
	'4':         "int",
	'5':         "int",
	'6':         "int",
	'7':         "int",
	'8':         "int",
	'9':         "int",
	TagInteger:  "int",
	TagLong:     "big.Int",
	TagDouble:   "float64",
	TagNull:     "nil",
	TagEmpty:    "empty string",
	TagTrue:     "true",
	TagFalse:    "false",
	TagNaN:      "NaN",
	TagInfinity: "Infinity",
	TagDate:     "time.Time",
	TagTime:     "time.Time",
	TagBytes:    "[]byte",
	TagUTF8Char: "string",
	TagString:   "string",
	TagGUID:     "GUID",
	TagList:     "slice",
	TagMap:      "map",
	TagClass:    "struct",
	TagObject:   "struct",
	TagRef:      "reference",
}

func tagToString(tag byte) (str string) {
	str = tagStringMap[tag]
	if str == "" {
		unexpectedTag(tag, nil)
	}
	return
}

func castError(tag byte, descType string) {
	srcType := tagToString(tag)
	panic(errors.New("can't convert " + srcType + " to " + descType))
}
<commit_msg>Added ReadTime method<commit_after>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * io/reader.go                                           *
 *                                                        *
 * hprose reader for Go.                                  *
 *                                                        *
 * LastModified: Oct 8, 2016                              *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package io

import (
	"bytes"
	"errors"
	"math/big"
	"reflect"
	"time"

	"github.com/hprose/hprose-golang/util"
)

// Reader is a fine-grained operation struct for Hprose unserialization
// when JSONCompatible is true, the Map data will unserialize to map[string]interface as the default type
type Reader struct {
	RawReader
	Simple         bool
	fieldsRef      [][]*fieldCache
	ref            []interface{}
	JSONCompatible bool
}

// NewReader is the constructor for Hprose Reader
func NewReader(buf []byte, simple bool) (reader *Reader) {
	reader = new(Reader)
	reader.buf = buf
	reader.Simple = simple
	return
}

// Unserialize a data from the reader
func (r *Reader) Unserialize(p interface{}) {
	v := reflect.ValueOf(p)
	if v.Kind() != reflect.Ptr {
		panic(errors.New("Unserialize: argument p must be a pointer"))
	}
	r.ReadValue(v.Elem())
}

// ReadValue from the reader
func (r *Reader) ReadValue(v reflect.Value) {
	tag := r.readByte()
	decoder := valueDecoders[v.Kind()]
	decoder(r, v, tag)
}

// CheckTag the next byte in reader is the expected tag or not
func (r *Reader) CheckTag(expectTag byte) (tag byte) {
	tag = r.readByte()
	if tag != expectTag {
		unexpectedTag(tag, []byte{expectTag})
	}
	return
}

// CheckTags the next byte in reader in the expected tags
func (r *Reader) CheckTags(expectTags []byte) (tag byte) {
	tag = r.readByte()
	if bytes.IndexByte(expectTags, tag) == -1 {
		unexpectedTag(tag, expectTags)
	}
	return
}

// ReadBool from the reader
func (r *Reader) ReadBool() bool {
	tag := r.readByte()
	decoder := boolDecoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "bool")
	return false
}

// ReadIntWithoutTag from the reader
func (r *Reader) ReadIntWithoutTag() int {
	return readInt(&r.ByteReader)
}

// ReadInt from the reader
func (r *Reader) ReadInt() int64 {
	tag := r.readByte()
	decoder := intDecoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "int64")
	return 0
}

// ReadUint from the reader
func (r *Reader) ReadUint() uint64 {
	tag := r.readByte()
	decoder := uintDecoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "uint64")
	return 0
}

// ReadFloat32 from the reader
func (r *Reader) ReadFloat32() float32 {
	tag := r.readByte()
	decoder := float32Decoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "float32")
	return 0
}

// ReadFloat64 from the reader
func (r *Reader) ReadFloat64() float64 {
	tag := r.readByte()
	decoder := float64Decoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "float64")
	return 0
}

// ReadComplex64 from the reader
func (r *Reader) ReadComplex64() complex64 {
	tag := r.readByte()
	decoder := complex64Decoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "complex64")
	return 0
}

// ReadComplex128 from the reader
func (r *Reader) ReadComplex128() complex128 {
	tag := r.readByte()
	decoder := complex128Decoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "complex128")
	return 0
}

// ReadStringWithoutTag from the reader
func (r *Reader) ReadStringWithoutTag() (str string) {
	str = readString(&r.ByteReader)
	if !r.Simple {
		setReaderRef(r, str)
	}
	return str
}

// ReadString from the reader
func (r *Reader) ReadString() (str string) {
	tag := r.readByte()
	decoder := stringDecoders[tag]
	if decoder != nil {
		return decoder(r)
	}
	castError(tag, "string")
	return ""
}

// ReadBytesWithoutTag from the reader
func (r *Reader) ReadBytesWithoutTag() (b []byte) {
	l := readLength(&r.ByteReader)
	b = make([]byte, l)
	if _, err := r.Read(b); err != nil {
		panic(err)
	}
	r.readByte()
	if !r.Simple {
		setReaderRef(r, b)
	}
	return
}

// ReadTime from the reader
func (r *Reader) ReadTime() (dt time.Time) {
	tag := r.readByte()
	switch tag {
	case TagDate:
		return r.ReadDateTimeWithoutTag()
	case TagTime:
		return r.ReadTimeWithoutTag()
	}
	castError(tag, "time.Time")
	return time.Time{}
}

// ReadDateTimeWithoutTag from the reader
func (r *Reader) ReadDateTimeWithoutTag() (dt time.Time) {
	year := r.read4Digit()
	month := r.read2Digit()
	day := r.read2Digit()
	tag := r.readByte()
	var hour, min, sec, nsec int
	if tag == TagTime {
		hour = r.read2Digit()
		min = r.read2Digit()
		sec = r.read2Digit()
		tag = r.readByte()
		if tag == TagPoint {
			nsec, tag = readNsec(&r.ByteReader)
		}
	}
	var loc *time.Location
	if tag == TagUTC {
		loc = time.UTC
	} else {
		loc = time.Local
	}
	dt = time.Date(year, time.Month(month), day, hour, min, sec, nsec, loc)
	if !r.Simple {
		setReaderRef(r, &dt)
	}
	return
}

// ReadTimeWithoutTag from the reader
func (r *Reader) ReadTimeWithoutTag() (t time.Time) {
	hour := r.read2Digit()
	min := r.read2Digit()
	sec := r.read2Digit()
	tag := r.readByte()
	var nsec int
	if tag == TagPoint {
		nsec, tag = readNsec(&r.ByteReader)
	}
	var loc *time.Location
	if tag == TagUTC {
		loc = time.UTC
	} else {
		loc = time.Local
	}
	t = time.Date(1970, 1, 1, hour, min, sec, nsec, loc)
	if !r.Simple {
		setReaderRef(r, &t)
	}
	return
}

// ReadBigIntWithoutTag from the reader
func (r *Reader) ReadBigIntWithoutTag() *big.Int {
	b := readUntil(&r.ByteReader, TagSemicolon)
	i, _ := new(big.Int).SetString(util.ByteString(b), 10)
	return i
}

// ReadSliceWithoutTag from the reader
func (r *Reader) ReadSliceWithoutTag() []reflect.Value {
	l := r.ReadCount()
	v := make([]reflect.Value, l)
	if !r.Simple {
		setReaderRef(r, nil)
	}
	for i := 0; i < l; i++ {
		v[i] = reflect.New(interfaceType).Elem()
		r.ReadValue(v[i])
	}
	r.readByte()
	return v
}

// ReadSlice from the reader
func (r *Reader) ReadSlice(v []reflect.Value) {
	l := len(v)
	if !r.Simple {
		setReaderRef(r, nil)
	}
	for i := 0; i < l; i++ {
		r.ReadValue(v[i])
	}
	r.readByte()
}

// ReadCount of array, slice, map or struct field
func (r *Reader) ReadCount() int {
	return int(ReadInt64(&r.ByteReader, TagOpenbrace))
}

// Reset the reference counter
func (r *Reader) Reset() {
	if r.fieldsRef != nil {
		r.fieldsRef = r.fieldsRef[:0]
	}
	if r.Simple {
		return
	}
	if r.ref != nil {
		r.ref = r.ref[:0]
	}
}

// private methods & functions

func (r *Reader) readRef() interface{} {
	if r.Simple {
		panic(errors.New("reference unserialization can't support in simple mode"))
	}
	return readRef(r, readInt(&r.ByteReader))
}

func setReaderRef(r *Reader, o interface{}) {
	r.ref = append(r.ref, o)
}

func readRef(r *Reader, i int) interface{} {
	return r.ref[i]
}

var tagStringMap = map[byte]string{
	'0':         "int",
	'1':         "int",
	'2':         "int",
	'3':         "int",
	'4':         "int",
	'5':         "int",
	'6':         "int",
	'7':         "int",
	'8':         "int",
	'9':         "int",
	TagInteger:  "int",
	TagLong:     "big.Int",
	TagDouble:   "float64",
	TagNull:     "nil",
	TagEmpty:    "empty string",
	TagTrue:     "true",
	TagFalse:    "false",
	TagNaN:      "NaN",
	TagInfinity: "Infinity",
	TagDate:     "time.Time",
	TagTime:     "time.Time",
	TagBytes:    "[]byte",
	TagUTF8Char: "string",
	TagString:   "string",
	TagGUID:     "GUID",
	TagList:     "slice",
	TagMap:      "map",
	TagClass:    "struct",
	TagObject:   "struct",
	TagRef:      "reference",
}

func tagToString(tag byte) (str string) {
	str = tagStringMap[tag]
	if str == "" {
		unexpectedTag(tag, nil)
	}
	return
}

func castError(tag byte, descType string) {
	srcType := tagToString(tag)
	panic(errors.New("can't convert " + srcType + " to " + descType))
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"socialapi/workers/common/manager"
	"socialapi/workers/common/runner"
	"socialapi/workers/pinnedpost/pinnedpost"
)

var (
	Name = "PinnedPost"
)

func main() {
	r := runner.New(Name)
	if err := r.Init(); err != nil {
		fmt.Println(err)
		return
	}

	// create message handler
	handler := pinnedpost.New(r.Log)

	m := manager.New()
	m.Controller(handler)

	m.HandleFunc("api.message_reply_created", (*pinnedpost.Controller).ReplyCreated)

	// create message handler
	r.Listen(m)
	r.Wait()
}
<commit_msg>Social: accept messages instead of message_replies<commit_after>package main

import (
	"fmt"
	"socialapi/workers/common/manager"
	"socialapi/workers/common/runner"
	"socialapi/workers/pinnedpost/pinnedpost"
)

var (
	Name = "PinnedPost"
)

func main() {
	r := runner.New(Name)
	if err := r.Init(); err != nil {
		fmt.Println(err)
		return
	}

	// create message handler
	handler := pinnedpost.New(r.Log)

	m := manager.New()
	m.Controller(handler)

	m.HandleFunc("api.channel_message_created", (*pinnedpost.Controller).MessageCreated)

	// create message handler
	r.Listen(m)
	r.Wait()
}
<|endoftext|>


<commit_before>package irc

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"strings"
)

type Client struct {
	conn        net.Conn
	User        User
	msgHandlers MessageHandlers
	Errorchan   chan error
}

func Connect(addr string, user User) (*Client, error) {
	client := &Client{
		User:      user,
		Errorchan: make(chan error),
	}
	client.setupMsgHandlers()
	cConn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, &Client{}
	}
	client.conn = cConn
	client.register(user)
	go client.handleInput()
	return client, nil
}

func (c *Client) Error() string {
	return "Error creating client.\n"
}

func (c *Client) Send(cmd string, a ...interface{}) {
	str := fmt.Sprintf(cmd, a...)
	c.conn.Write([]byte(str + "\r\n"))
	log.Println("out>", str)
}

func (c *Client) Join(channel, password string) {
	c.Send("JOIN %s %s", channel, password)
}

func (c *Client) Nick(nick string) {
	c.Send("NICK " + nick)
}

func (c *Client) Notice(to, msg string) {
	c.Send("NOTICE %s :%s", to, msg)
}

func (c *Client) Part(channel string) {
	c.Send("PART " + channel)
}

func (c *Client) Ping(arg string) {
	c.Send("PING :" + arg)
}

func (c *Client) Pong(arg string) {
	c.Send("PONG :" + arg)
}

func (c *Client) PrivMsg(to, msg string) {
	c.Send("PRIVMSG %s :%s", to, msg)
}

func (c *Client) register(user User) {
	c.Nick(user.Nick)
	c.Send("USER %s %d * :%s", user.Nick, user.mode, user.Realname)
}

func (c *Client) responseCTCP(to, answer string) {
	c.Notice(to, ctcpQuote(answer))
}

func (c *Client) respondTo(maskedUser, action, talkedTo, message string) {
	maskedUser = strings.TrimPrefix(maskedUser, ":")
	message = strings.TrimPrefix(message, ":")
	user := strings.SplitN(maskedUser, "!", 2)
	msg := &Message{
		From:   user[0],
		To:     talkedTo,
		Action: action,
		Body:   message,
	}
	for _, fn := range c.msgHandlers[action] {
		go fn(msg)
	}
}

func (c *Client) handleInput() {
	defer c.conn.Close()
	reader := bufio.NewReader(c.conn)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			c.Errorchan <- err
			break
		}
		// FIXME: This is obviously not the right way to parse messages
		packet := strings.SplitN(line[:len(line)-2], " ", 4)
		if len(packet) == 2 {
			packet = []string{packet[1], packet[0], c.User.Nick, packet[1]}
		}
		if len(packet) == 4 {
			go c.respondTo(packet[0], packet[1], packet[2], packet[3])
		}
		log.Println("in>", line)
	}

}
<commit_msg>Use fixed number of goroutines, use more channels.<commit_after>package irc

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"strings"
)

type Client struct {
	conn        net.Conn
	User        User
	msgHandlers MessageHandlers
	Errorchan   chan error
	messagechan chan *Message
}

func Connect(addr string, user User) (*Client, error) {
	client := &Client{
		User:      user,
		Errorchan: make(chan error),
		// Buffer 25 messages at channel, this is a lot!
		messagechan: make(chan *Message, 25),
	}
	client.setupMsgHandlers()
	cConn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, &Client{}
	}
	client.conn = cConn
	client.register(user)
	go client.handleInput()
	go client.processMessage()
	return client, nil
}

func (c *Client) Error() string {
	return "Error creating client.\n"
}

func (c *Client) Send(cmd string, a ...interface{}) {
	str := fmt.Sprintf(cmd, a...)
	c.conn.Write([]byte(str + "\r\n"))
	log.Println("out>", str)
}

func (c *Client) Join(channel, password string) {
	c.Send("JOIN %s %s", channel, password)
}

func (c *Client) Nick(nick string) {
	c.Send("NICK " + nick)
}

func (c *Client) Notice(to, msg string) {
	c.Send("NOTICE %s :%s", to, msg)
}

func (c *Client) Part(channel string) {
	c.Send("PART " + channel)
}

func (c *Client) Ping(arg string) {
	c.Send("PING :" + arg)
}

func (c *Client) Pong(arg string) {
	c.Send("PONG :" + arg)
}

func (c *Client) PrivMsg(to, msg string) {
	c.Send("PRIVMSG %s :%s", to, msg)
}

func (c *Client) register(user User) {
	c.Nick(user.Nick)
	c.Send("USER %s %d * :%s", user.Nick, user.mode, user.Realname)
	if len(user.password) != 0 {
		c.PrivMsg("nickserv", "identify "+user.password)
	}
}

func (c *Client) responseCTCP(to, answer string) {
	c.Notice(to, ctcpQuote(answer))
}

func (c *Client) respondTo(maskedUser, action, talkedTo, message string) {
	message = strings.TrimPrefix(message, ":")
	maskedUser = strings.TrimPrefix(maskedUser, ":")
	user := strings.SplitN(maskedUser, "!", 2)
	msg := &Message{
		From:   user[0],
		To:     talkedTo,
		Action: action,
		Body:   message,
	}
	c.messagechan <- msg
}

func (c *Client) handleInput() {
	defer c.conn.Close()
	reader := bufio.NewReader(c.conn)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			close(c.messagechan)
			c.Errorchan <- err
			break
		}
		log.Println("in>", line)
		// FIXME: This is obviously not the right way to parse messages
		packet := strings.SplitN(line[:len(line)-2], " ", 4)
		if len(packet) == 2 {
			packet = []string{packet[1], packet[0], c.User.Nick, packet[1]}
		}
		if len(packet) == 4 {
			c.respondTo(packet[0], packet[1], packet[2], packet[3])
		}
	}
}

func (c *Client) processMessage() {
	for {
		msg, ok := <-c.messagechan
		if !ok {
			return
		}
		for _, fn := range c.msgHandlers[msg.Action] {
			fn(msg)
		}
	}
}
<|endoftext|>


<commit_before>package nacl

// #include <sodium.h>
// #cgo pkg-config: libsodium
import "C"

import (
  "bytes"
  "log"
)

// return how many bytes overhead does CryptoBox have
func CryptoBoxOverhead() int {
  return int(C.crypto_box_macbytes())
}

// size of crypto_box public keys
func CryptoBoxPubKeySize() int {
  return int(C.crypto_box_publickeybytes())
}

// size of crypto_box private keys
func CryptoBoxPrivKeySize() int {
  return int(C.crypto_box_secretkeybytes())
}

// size of crypto_sign public keys
func CryptoSignPubKeySize() int {
  return int(C.crypto_sign_publickeybytes())
}

// size of crypto_sign private keys
func CryptoSignPrivKeySize() int {
  return int(C.crypto_sign_secretkeybytes())
}

func testSign(bufflen int, keys *KeyPair) {
  log.Printf("Test %d sign/verify...", bufflen)
  msg := RandBytes(bufflen)
  sig := CryptoSignDetached(msg, keys.sk.Data())
  if ! CryptoVerifyDetached(msg, sig, keys.pk.Data()) {
    log.Fatal("Failed")
  }
}

func testBox(bufflen int, tokey, fromkey *KeyPair) {
  log.Printf("Test %d box/box_open...", bufflen)
  msg := RandBytes(bufflen)
  nounce := NewBoxNounce()
  box := CryptoBox(msg, nounce, tokey.Public(), fromkey.Secret())
  if box == nil {
    log.Fatal("CryptoBox() failed")
  }
  msg_open := CryptoBoxOpen(box, nounce, tokey.Secret(), fromkey.Public())
  if ! bytes.Equal(msg, msg_open) {
    log.Fatalf("CryptoBoxOpen() failed: %d vs %d", len(msg), len(msg_open))
  }
}

// test all crypto functions
func TestAll() {
  log.Println("Begin Crypto test")

  bufflen := 128

  b := RandBytes(bufflen)
  if len(b) != bufflen {
    log.Fatal("nacl.RandBytes() failed length test")
  }

  for n := 1 ; n < 16 ; n++ {
    key := GenSignKeypair()
    defer key.Free()
    testSign(n * 1024, key)
  }

  for n := 1 ; n < 16 ; n++ {
    tokey := GenBoxKeypair()
    fromkey := GenBoxKeypair()
    defer tokey.Free()
    defer fromkey.Free()
    testBox(n * 1024, tokey, fromkey)
  }

  log.Println("Crypto Test Done")
}

// initialize sodium
func init() {
  status := C.sodium_init()
  if status == -1 {
    log.Fatalf("failed to initialize libsodium status=%d", status)
  }
  version_ptr := C.sodium_version_string()

  log.Println("Intialized Sodium", C.GoString(version_ptr))
  TestAll()
}
<commit_msg>crypto test now optional<commit_after>package nacl

// #include <sodium.h>
// #cgo pkg-config: libsodium
import "C"

import (
  "bytes"
  "log"
  "os"
)

// return how many bytes overhead does CryptoBox have
func CryptoBoxOverhead() int {
  return int(C.crypto_box_macbytes())
}

// size of crypto_box public keys
func CryptoBoxPubKeySize() int {
  return int(C.crypto_box_publickeybytes())
}

// size of crypto_box private keys
func CryptoBoxPrivKeySize() int {
  return int(C.crypto_box_secretkeybytes())
}

// size of crypto_sign public keys
func CryptoSignPubKeySize() int {
  return int(C.crypto_sign_publickeybytes())
}

// size of crypto_sign private keys
func CryptoSignPrivKeySize() int {
  return int(C.crypto_sign_secretkeybytes())
}

func testSign(bufflen int, keys *KeyPair) {
  log.Printf("Test %d sign/verify...", bufflen)
  msg := RandBytes(bufflen)
  sig := CryptoSignDetached(msg, keys.sk.Data())
  if ! CryptoVerifyDetached(msg, sig, keys.pk.Data()) {
    log.Fatal("Failed")
  }
}

func testBox(bufflen int, tokey, fromkey *KeyPair) {
  log.Printf("Test %d box/box_open...", bufflen)
  msg := RandBytes(bufflen)
  nounce := NewBoxNounce()
  box := CryptoBox(msg, nounce, tokey.Public(), fromkey.Secret())
  if box == nil {
    log.Fatal("CryptoBox() failed")
  }
  msg_open := CryptoBoxOpen(box, nounce, tokey.Secret(), fromkey.Public())
  if ! bytes.Equal(msg, msg_open) {
    log.Fatalf("CryptoBoxOpen() failed: %d vs %d", len(msg), len(msg_open))
  }
}

// test all crypto functions
func TestAll() {
  log.Println("Begin Crypto test")

  bufflen := 128

  b := RandBytes(bufflen)
  if len(b) != bufflen {
    log.Fatal("nacl.RandBytes() failed length test")
  }

  for n := 1 ; n < 16 ; n++ {
    key := GenSignKeypair()
    defer key.Free()
    testSign(n * 1024, key)
  }

  for n := 1 ; n < 16 ; n++ {
    tokey := GenBoxKeypair()
    fromkey := GenBoxKeypair()
    defer tokey.Free()
    defer fromkey.Free()
    testBox(n * 1024, tokey, fromkey)
  }

  log.Println("Crypto Test Done")
}

// initialize sodium
func init() {
  status := C.sodium_init()
  if status == -1 {
    log.Fatalf("failed to initialize libsodium status=%d", status)
  }

  if os.Getenv("SODIUM_TEST") == "1" {
    version_ptr := C.sodium_version_string()
    log.Println("Intialized Sodium", C.GoString(version_ptr))
    TestAll()
  }
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/go-martini/martini"
	"github.com/gorilla/websocket"
	"github.com/martini-contrib/sessions"
	"log"
	"math/rand"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"
)

const CLEAR string = "\\33[H\\33[2J"

type Config struct {
	Secret  string
	Address string
}

type WebsocketData struct {
	T string
	D []byte
}

type LockingWebsockets struct {
	mutex sync.RWMutex
	byId  map[int64]*websocket.Conn
}

var websockets *LockingWebsockets

func (c *LockingWebsockets) deleteWebsocket(id int64) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	delete(c.byId, id)
}

func (c *LockingWebsockets) addWebsocket(id int64, ws *websocket.Conn) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	c.byId[id] = ws
}

func main() {
	websockets = &LockingWebsockets{
		byId: make(map[int64]*websocket.Conn),
	}
	consoleReadChannel = make(chan ConsoleChunk)
	go consoleDispatch()
	vzcontrol := ConnectVZControl()
	defer vzcontrol.Close()

	file, _ := os.Open("game_server.cfg")
	decoder := json.NewDecoder(file)
	config := Config{}
	decoder.Decode(&config)

	m := martini.Classic()
	store := sessions.NewCookieStore([]byte(config.Secret))
	m.Use(sessions.Sessions("session", store))

	generating := false
	gr := NewGraph()
	m.Get("/reset/:secret", func(w http.ResponseWriter, r *http.Request, params martini.Params, session sessions.Session) string {
		if params["secret"] != config.Secret {
			return ""
		}
		err := vzcontrol.Reset()
		if err != nil {
			return err.Error()
		}
		generating = false
		gr = NewGraph()
		return "Done"
	})

	m.Get("/gen", func(w http.ResponseWriter, r *http.Request, session sessions.Session) string {
		if generating {
			return "Already generating"
		}
		generating = true
		maxNodes := 5
		maxEdges := 5
		startNodeId := 100

		startNode := Node{Id: NodeId(startNodeId)}
		gr.AddNode(startNode)
		err := vzcontrol.ContainerCreate(int64(startNode.Id))

		nodes := make([]Node, 0)
		nodes = append(nodes, startNode)

		steps := 1
		for len(nodes) != 0 && steps < maxNodes {
			node, nodes := nodes[len(nodes)-1], nodes[:len(nodes)-1]

			numEdges := random(1, maxEdges)
			for i := 1; i <= numEdges; i++ {
				targetNode := Node{Id: NodeId(i*steps + startNodeId)}
				if gr.AddNode(targetNode) {
					err = vzcontrol.ContainerCreate(int64(targetNode.Id))
					if err != nil {
						return fmt.Sprintf("Container Create: %d, %d, %d\n%s", targetNode.Id, i*steps, numEdges, err.Error())
					}
					nodes = append(nodes, targetNode)
					edgeid := int64(i * steps)
					if gr.AddEdge(Edge{Id: EdgeId(edgeid), Head: node.Id, Tail: targetNode.Id}) {
						err = vzcontrol.NetworkCreate(edgeid)
						if err != nil {
							return fmt.Sprintf("Network Create: %d\n%s", edgeid, err.Error())
						}
						err = vzcontrol.NetworkAdd(int64(node.Id), edgeid)
						if err != nil {
							return fmt.Sprintf("Network Add Node: %d, %d\n%s", node.Id, edgeid, err.Error())
						}
						err = vzcontrol.NetworkAdd(int64(targetNode.Id), edgeid)
						if err != nil {
							return fmt.Sprintf("Network Add Target: %d, %d\n%s", targetNode.Id, edgeid, err.Error())
						}
					}
				}

			}
			steps += 1
		}
		return gr.String()
	})

	m.Get("/graph", func(w http.ResponseWriter, r *http.Request, session sessions.Session) string {
		output, err := json.Marshal(gr)
		if err != nil {
			return err.Error()
		}
		return string(output)
	})

	m.Get("/ws", func(w http.ResponseWriter, r *http.Request, session sessions.Session) {
		currentVm := 0
		ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		if _, ok := err.(websocket.HandshakeError); ok {
			http.Error(w, "Not a websocket handshake", 400)
			log.Println(err)
			return
		} else if err != nil {
			log.Println(err)
			return
		}
		defer ws.Close()
		ws.WriteMessage(websocket.TextMessage, []byte("Welcome to ginux!\r\n"))
		for {
			_, message, err := ws.ReadMessage()
			if err != nil {
				log.Println(err)
				return
			} else {
				msg := WebsocketData{}
				json.Unmarshal(message, &msg)
				switch msg.T {
				case "term":
					if currentVm != 0 {
						vzcontrol.ConsoleWrite(int64(currentVm), msg.D)
					}
				case "click":
					currentVm, _ = strconv.Atoi(string(msg.D))
					ws.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf("%d", currentVm)))
				}
			}
		}
	})
	log.Println("Game Server started on", config.Address)
	log.Fatal(http.ListenAndServe(config.Address, m))
}

func random(min, max int) int {
	rand.Seed(time.Now().Unix())
	return rand.Intn(max-min) + min
}

func consoleDispatch() {
	for chunk := range consoleReadChannel {
		websockets.mutex.RLock()
		if socket, ok := websockets.byId[chunk.Id]; ok {
			socket.WriteMessage(websocket.TextMessage, chunk.Data)
		}
		websockets.mutex.RUnlock()
	}
}
<commit_msg>Data zero<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/go-martini/martini"
	"github.com/gorilla/websocket"
	"github.com/martini-contrib/sessions"
	"log"
	"math/rand"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"
)

const CLEAR string = "\\33[H\\33[2J"

type Config struct {
	Secret  string
	Address string
}

type WebsocketData struct {
	T string
	D []byte
}

type LockingWebsockets struct {
	mutex sync.RWMutex
	byId  map[int64]*websocket.Conn
}

var websockets *LockingWebsockets

func (c *LockingWebsockets) deleteWebsocket(id int64) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	delete(c.byId, id)
}

func (c *LockingWebsockets) addWebsocket(id int64, ws *websocket.Conn) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	c.byId[id] = ws
}

func main() {
	websockets = &LockingWebsockets{
		byId: make(map[int64]*websocket.Conn),
	}
	consoleReadChannel = make(chan ConsoleChunk)
	go consoleDispatch()
	vzcontrol := ConnectVZControl()
	defer vzcontrol.Close()

	file, _ := os.Open("game_server.cfg")
	decoder := json.NewDecoder(file)
	config := Config{}
	decoder.Decode(&config)

	m := martini.Classic()
	store := sessions.NewCookieStore([]byte(config.Secret))
	m.Use(sessions.Sessions("session", store))

	generating := false
	gr := NewGraph()
	m.Get("/reset/:secret", func(w http.ResponseWriter, r *http.Request, params martini.Params, session sessions.Session) string {
		if params["secret"] != config.Secret {
			return ""
		}
		err := vzcontrol.Reset()
		if err != nil {
			return err.Error()
		}
		generating = false
		gr = NewGraph()
		return "Done"
	})

	m.Get("/gen", func(w http.ResponseWriter, r *http.Request, session sessions.Session) string {
		if generating {
			return "Already generating"
		}
		generating = true
		maxNodes := 5
		maxEdges := 5
		startNodeId := 100

		startNode := Node{Id: NodeId(startNodeId)}
		gr.AddNode(startNode)
		err := vzcontrol.ContainerCreate(int64(startNode.Id))

		nodes := make([]Node, 0)
		nodes = append(nodes, startNode)

		steps := 1
		for len(nodes) != 0 && steps < maxNodes {
			node, nodes := nodes[len(nodes)-1], nodes[:len(nodes)-1]

			numEdges := random(1, maxEdges)
			for i := 1; i <= numEdges; i++ {
				targetNode := Node{Id: NodeId(i*steps + startNodeId)}
				if gr.AddNode(targetNode) {
					err = vzcontrol.ContainerCreate(int64(targetNode.Id))
					if err != nil {
						return fmt.Sprintf("Container Create: %d, %d, %d\n%s", targetNode.Id, i*steps, numEdges, err.Error())
					}
					nodes = append(nodes, targetNode)
					edgeid := int64(i * steps)
					if gr.AddEdge(Edge{Id: EdgeId(edgeid), Head: node.Id, Tail: targetNode.Id}) {
						err = vzcontrol.NetworkCreate(edgeid)
						if err != nil {
							return fmt.Sprintf("Network Create: %d\n%s", edgeid, err.Error())
						}
						err = vzcontrol.NetworkAdd(int64(node.Id), edgeid)
						if err != nil {
							return fmt.Sprintf("Network Add Node: %d, %d\n%s", node.Id, edgeid, err.Error())
						}
						err = vzcontrol.NetworkAdd(int64(targetNode.Id), edgeid)
						if err != nil {
							return fmt.Sprintf("Network Add Target: %d, %d\n%s", targetNode.Id, edgeid, err.Error())
						}
					}
				}

			}
			steps += 1
		}
		return gr.String()
	})

	m.Get("/graph", func(w http.ResponseWriter, r *http.Request, session sessions.Session) string {
		output, err := json.Marshal(gr)
		if err != nil {
			return err.Error()
		}
		return string(output)
	})

	m.Get("/ws", func(w http.ResponseWriter, r *http.Request, session sessions.Session) {
		currentVm := 0
		ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		if _, ok := err.(websocket.HandshakeError); ok {
			http.Error(w, "Not a websocket handshake", 400)
			log.Println(err)
			return
		} else if err != nil {
			log.Println(err)
			return
		}
		defer ws.Close()
		ws.WriteMessage(websocket.TextMessage, []byte("Welcome to ginux!\r\n"))
		for {
			_, message, err := ws.ReadMessage()
			if err != nil {
				log.Println(err)
				return
			} else {
				msg := WebsocketData{}
				json.Unmarshal(message, &msg)
				switch msg.T {
				case "term":
					if currentVm != 0 {
						vzcontrol.ConsoleWrite(int64(currentVm), msg.D)
					}
				case "click":
					ws.WriteMessage(websocket.TextMessage, msg.D)
				}
			}
		}
	})
	log.Println("Game Server started on", config.Address)
	log.Fatal(http.ListenAndServe(config.Address, m))
}

func random(min, max int) int {
	rand.Seed(time.Now().Unix())
	return rand.Intn(max-min) + min
}

func consoleDispatch() {
	for chunk := range consoleReadChannel {
		websockets.mutex.RLock()
		if socket, ok := websockets.byId[chunk.Id]; ok {
			socket.WriteMessage(websocket.TextMessage, chunk.Data)
		}
		websockets.mutex.RUnlock()
	}
}
<|endoftext|>


<commit_before>package handlers

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"strings"

	"github.com/docker/notary/tuf/data"
	"github.com/docker/notary/tuf/signed"
	"github.com/gorilla/mux"
	"golang.org/x/net/context"

	ctxu "github.com/docker/distribution/context"
	"github.com/docker/notary/errors"
	"github.com/docker/notary/server/storage"
	"github.com/docker/notary/server/timestamp"
)

// MainHandler is the default handler for the server
func MainHandler(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	if r.Method == "GET" {
		_, err := w.Write([]byte("{}"))
		if err != nil {
			return errors.ErrUnknown.WithDetail(err)
		}
	} else {
		return errors.ErrGenericNotFound.WithDetail(nil)
	}
	return nil
}

// AtomicUpdateHandler will accept multiple TUF files and ensure that the storage
// backend is atomically updated with all the new records.
func AtomicUpdateHandler(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	defer r.Body.Close()
	s := ctx.Value("metaStore")
	store, ok := s.(storage.MetaStore)
	if !ok {
		return errors.ErrNoStorage.WithDetail(nil)
	}
	cryptoServiceVal := ctx.Value("cryptoService")
	cryptoService, ok := cryptoServiceVal.(signed.CryptoService)
	if !ok {
		return errors.ErrNoCryptoService.WithDetail(nil)
	}

	vars := mux.Vars(r)
	gun := vars["imageName"]
	reader, err := r.MultipartReader()
	if err != nil {
		return errors.ErrMalformedUpload.WithDetail(nil)
	}
	var updates []storage.MetaUpdate
	for {
		part, err := reader.NextPart()
		if err == io.EOF {
			break
		}
		role := strings.TrimSuffix(part.FileName(), ".json")
		if role == "" {
			return errors.ErrNoFilename.WithDetail(nil)
		} else if !data.ValidRole(role) {
			return errors.ErrInvalidRole.WithDetail(role)
		}
		meta := &data.SignedMeta{}
		var input []byte
		inBuf := bytes.NewBuffer(input)
		dec := json.NewDecoder(io.TeeReader(part, inBuf))
		err = dec.Decode(meta)
		if err != nil {
			return errors.ErrMalformedJSON.WithDetail(nil)
		}
		version := meta.Signed.Version
		updates = append(updates, storage.MetaUpdate{
			Role:    role,
			Version: version,
			Data:    inBuf.Bytes(),
		})
	}
	updates, err = validateUpdate(cryptoService, gun, updates, store)
	if err != nil {
		return errors.ErrMalformedUpload.WithDetail(err)
	}
	err = store.UpdateMany(gun, updates)
	if err != nil {
		return errors.ErrUpdating.WithDetail(err)
	}
	return nil
}

// GetHandler returns the json for a specified role and GUN.
func GetHandler(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	s := ctx.Value("metaStore")
	store, ok := s.(storage.MetaStore)
	if !ok {
		return errors.ErrNoStorage.WithDetail(nil)
	}
	vars := mux.Vars(r)
	gun := vars["imageName"]
	tufRole := vars["tufRole"]

	logger := ctxu.GetLoggerWithFields(ctx, map[string]interface{}{"gun": gun, "tufRole": tufRole})

	if data.CanonicalRole(tufRole) == data.CanonicalTimestampRole {
		return getTimestamp(ctx, w, logger, store, gun)
	}

	out, err := store.GetCurrent(gun, tufRole)
	if err != nil {
		if _, ok := err.(*storage.ErrNotFound); ok {
			return errors.ErrMetadataNotFound.WithDetail(nil)
		}
		logger.Error("500 GET")
		return errors.ErrUnknown.WithDetail(err)
	}
	if out == nil {
		logger.Error("404 GET")
		return errors.ErrMetadataNotFound.WithDetail(nil)
	}
	w.Write(out)
	logger.Debug("200 GET")

	return nil
}

// DeleteHandler deletes all data for a GUN. A 200 responses indicates success.
func DeleteHandler(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	s := ctx.Value("metaStore")
	store, ok := s.(storage.MetaStore)
	if !ok {
		return errors.ErrNoStorage.WithDetail(nil)
	}
	vars := mux.Vars(r)
	gun := vars["imageName"]
	logger := ctxu.GetLoggerWithField(ctx, gun, "gun")
	err := store.Delete(gun)
	if err != nil {
		logger.Error("500 DELETE repository")
		return errors.ErrUnknown.WithDetail(err)
	}
	return nil
}

// getTimestampHandler returns a timestamp.json given a GUN
func getTimestamp(ctx context.Context, w http.ResponseWriter, logger ctxu.Logger, store storage.MetaStore, gun string) error {
	cryptoServiceVal := ctx.Value("cryptoService")
	cryptoService, ok := cryptoServiceVal.(signed.CryptoService)
	if !ok {
		return errors.ErrNoCryptoService.WithDetail(nil)
	}

	out, err := timestamp.GetOrCreateTimestamp(gun, store, cryptoService)
	if err != nil {
		switch err.(type) {
		case *storage.ErrNoKey, *storage.ErrNotFound:
			logger.Error("404 GET timestamp")
			return errors.ErrMetadataNotFound.WithDetail(nil)
		default:
			logger.Error("500 GET timestamp")
			return errors.ErrUnknown.WithDetail(err)
		}
	}

	logger.Debug("200 GET timestamp")
	w.Write(out)
	return nil
}

// GetTimestampKeyHandler returns a timestamp public key, creating a new key-pair
// it if it doesn't yet exist
func GetTimestampKeyHandler(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	vars := mux.Vars(r)
	gun := vars["imageName"]

	logger := ctxu.GetLoggerWithField(ctx, gun, "gun")

	s := ctx.Value("metaStore")
	store, ok := s.(storage.MetaStore)
	if !ok {
		logger.Error("500 GET storage not configured")
		return errors.ErrNoStorage.WithDetail(nil)
	}
	c := ctx.Value("cryptoService")
	crypto, ok := c.(signed.CryptoService)
	if !ok {
		logger.Error("500 GET crypto service not configured")
		return errors.ErrNoCryptoService.WithDetail(nil)
	}
	algo := ctx.Value("keyAlgorithm")
	keyAlgo, ok := algo.(string)
	if !ok {
		logger.Error("500 GET key algorithm not configured")
		return errors.ErrNoKeyAlgorithm.WithDetail(nil)
	}
	keyAlgorithm := keyAlgo

	key, err := timestamp.GetOrCreateTimestampKey(gun, store, crypto, keyAlgorithm)
	if err != nil {
		logger.Errorf("500 GET timestamp key: %v", err)
		return errors.ErrUnknown.WithDetail(err)
	}

	out, err := json.Marshal(key)
	if err != nil {
		logger.Error("500 GET timestamp key")
		return errors.ErrUnknown.WithDetail(err)
	}
	logger.Debug("200 GET timestamp key")
	w.Write(out)
	return nil
}
<commit_msg>update GetHandler to use new snapshot code Signed-off-by: David Lawrence <david.lawrence@docker.com> (github: endophage)<commit_after>package handlers

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"strings"

	ctxu "github.com/docker/distribution/context"
	"github.com/gorilla/mux"
	"golang.org/x/net/context"

	"github.com/Sirupsen/logrus"
	"github.com/docker/notary/errors"
	"github.com/docker/notary/server/snapshot"
	"github.com/docker/notary/server/storage"
	"github.com/docker/notary/server/timestamp"
	"github.com/docker/notary/tuf/data"
	"github.com/docker/notary/tuf/signed"
)

// MainHandler is the default handler for the server
func MainHandler(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	if r.Method == "GET" {
		_, err := w.Write([]byte("{}"))
		if err != nil {
			return errors.ErrUnknown.WithDetail(err)
		}
	} else {
		return errors.ErrGenericNotFound.WithDetail(nil)
	}
	return nil
}

// AtomicUpdateHandler will accept multiple TUF files and ensure that the storage
// backend is atomically updated with all the new records.
func AtomicUpdateHandler(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	defer r.Body.Close()
	s := ctx.Value("metaStore")
	store, ok := s.(storage.MetaStore)
	if !ok {
		return errors.ErrNoStorage.WithDetail(nil)
	}
	cryptoServiceVal := ctx.Value("cryptoService")
	cryptoService, ok := cryptoServiceVal.(signed.CryptoService)
	if !ok {
		return errors.ErrNoCryptoService.WithDetail(nil)
	}

	vars := mux.Vars(r)
	gun := vars["imageName"]
	reader, err := r.MultipartReader()
	if err != nil {
		return errors.ErrMalformedUpload.WithDetail(nil)
	}
	var updates []storage.MetaUpdate
	for {
		part, err := reader.NextPart()
		if err == io.EOF {
			break
		}
		role := strings.TrimSuffix(part.FileName(), ".json")
		if role == "" {
			return errors.ErrNoFilename.WithDetail(nil)
		} else if !data.ValidRole(role) {
			return errors.ErrInvalidRole.WithDetail(role)
		}
		meta := &data.SignedMeta{}
		var input []byte
		inBuf := bytes.NewBuffer(input)
		dec := json.NewDecoder(io.TeeReader(part, inBuf))
		err = dec.Decode(meta)
		if err != nil {
			return errors.ErrMalformedJSON.WithDetail(nil)
		}
		version := meta.Signed.Version
		updates = append(updates, storage.MetaUpdate{
			Role:    role,
			Version: version,
			Data:    inBuf.Bytes(),
		})
	}
	updates, err = validateUpdate(cryptoService, gun, updates, store)
	if err != nil {
		return errors.ErrMalformedUpload.WithDetail(err)
	}
	err = store.UpdateMany(gun, updates)
	if err != nil {
		return errors.ErrUpdating.WithDetail(err)
	}
	return nil
}

// GetHandler returns the json for a specified role and GUN.
func GetHandler(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	s := ctx.Value("metaStore")
	store, ok := s.(storage.MetaStore)
	if !ok {
		return errors.ErrNoStorage.WithDetail(nil)
	}
	vars := mux.Vars(r)
	gun := vars["imageName"]
	tufRole := vars["tufRole"]

	logger := ctxu.GetLoggerWithFields(ctx, map[string]interface{}{"gun": gun, "tufRole": tufRole})

	switch data.CanonicalRole(tufRole) {
	case data.CanonicalTimestampRole:
		return getTimestamp(ctx, w, logger, store, gun)
	case data.CanonicalSnapshotRole:
		return getSnapshot(ctx, w, logger, store, gun)
	}

	out, err := store.GetCurrent(gun, tufRole)
	if err != nil {
		if _, ok := err.(*storage.ErrNotFound); ok {
			logrus.Error(gun + ":" + tufRole)
			return errors.ErrMetadataNotFound.WithDetail(nil)
		}
		logger.Error("500 GET")
		return errors.ErrUnknown.WithDetail(err)
	}
	if out == nil {
		logger.Error("404 GET")
		return errors.ErrMetadataNotFound.WithDetail(nil)
	}
	w.Write(out)
	logger.Debug("200 GET")

	return nil
}

// DeleteHandler deletes all data for a GUN. A 200 responses indicates success.
func DeleteHandler(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	s := ctx.Value("metaStore")
	store, ok := s.(storage.MetaStore)
	if !ok {
		return errors.ErrNoStorage.WithDetail(nil)
	}
	vars := mux.Vars(r)
	gun := vars["imageName"]
	logger := ctxu.GetLoggerWithField(ctx, gun, "gun")
	err := store.Delete(gun)
	if err != nil {
		logger.Error("500 DELETE repository")
		return errors.ErrUnknown.WithDetail(err)
	}
	return nil
}

// getTimestampHandler returns a timestamp.json given a GUN
func getTimestamp(ctx context.Context, w http.ResponseWriter, logger ctxu.Logger, store storage.MetaStore, gun string) error {
	cryptoServiceVal := ctx.Value("cryptoService")
	cryptoService, ok := cryptoServiceVal.(signed.CryptoService)
	if !ok {
		return errors.ErrNoCryptoService.WithDetail(nil)
	}

	out, err := timestamp.GetOrCreateTimestamp(gun, store, cryptoService)
	if err != nil {
		switch err.(type) {
		case *storage.ErrNoKey, *storage.ErrNotFound:
			logger.Error("404 GET timestamp")
			return errors.ErrMetadataNotFound.WithDetail(nil)
		default:
			logger.Error("500 GET timestamp")
			return errors.ErrUnknown.WithDetail(err)
		}
	}

	logger.Debug("200 GET timestamp")
	w.Write(out)
	return nil
}

// getTimestampHandler returns a timestamp.json given a GUN
func getSnapshot(ctx context.Context, w http.ResponseWriter, logger ctxu.Logger, store storage.MetaStore, gun string) error {
	cryptoServiceVal := ctx.Value("cryptoService")
	cryptoService, ok := cryptoServiceVal.(signed.CryptoService)
	if !ok {
		return errors.ErrNoCryptoService.WithDetail(nil)
	}

	out, err := snapshot.GetOrCreateSnapshot(gun, store, cryptoService)
	if err != nil {
		switch err.(type) {
		case *storage.ErrNoKey, *storage.ErrNotFound:
			logger.Error("404 GET snapshot")
			return errors.ErrMetadataNotFound.WithDetail(nil)
		default:
			logger.Error("500 GET snapshot")
			return errors.ErrUnknown.WithDetail(err)
		}
	}

	logger.Debug("200 GET snapshot")
	w.Write(out)
	return nil
}

// GetTimestampKeyHandler returns a timestamp public key, creating a new key-pair
// it if it doesn't yet exist
func GetTimestampKeyHandler(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	vars := mux.Vars(r)
	gun := vars["imageName"]

	logger := ctxu.GetLoggerWithField(ctx, gun, "gun")

	s := ctx.Value("metaStore")
	store, ok := s.(storage.MetaStore)
	if !ok {
		logger.Error("500 GET storage not configured")
		return errors.ErrNoStorage.WithDetail(nil)
	}
	c := ctx.Value("cryptoService")
	crypto, ok := c.(signed.CryptoService)
	if !ok {
		logger.Error("500 GET crypto service not configured")
		return errors.ErrNoCryptoService.WithDetail(nil)
	}
	algo := ctx.Value("keyAlgorithm")
	keyAlgo, ok := algo.(string)
	if !ok {
		logger.Error("500 GET key algorithm not configured")
		return errors.ErrNoKeyAlgorithm.WithDetail(nil)
	}
	keyAlgorithm := keyAlgo

	key, err := timestamp.GetOrCreateTimestampKey(gun, store, crypto, keyAlgorithm)
	if err != nil {
		logger.Errorf("500 GET timestamp key: %v", err)
		return errors.ErrUnknown.WithDetail(err)
	}

	out, err := json.Marshal(key)
	if err != nil {
		logger.Error("500 GET timestamp key")
		return errors.ErrUnknown.WithDetail(err)
	}
	logger.Debug("200 GET timestamp key")
	w.Write(out)
	return nil
}
<|endoftext|>


<commit_before>package storage

import (
	"github.com/Sirupsen/logrus"
	"github.com/endophage/gotuf/data"
	"github.com/jinzhu/gorm"
)

// SQLStorage implements a versioned store using a relational database.
// The database table must look like:
// CREATE TABLE `tuf_files` (
//   `id` INT AUTO_INCREMENT,
//   `gun` VARCHAR(255) NOT NULL
//   `role` VARCHAR(255) NOT NULL
//   `version` INT
//   `data` LONGBLOB
//   PRIMARY KEY (`id`)
//   UNIQUE INDEX (`gun`, `role`, `version`)
// ) DEFAULT CHARSET=utf8;
//
// CREATE TABLE `timestamp_keys` (
//   `gun` VARCHAR(255),
//   `cipher` VARCHAR(30),
//   `public` BLOB NOT NULL,
//    PRIMARY KEY (`gun`)
// ) DEFAULT CHARSET=utf8;
type SQLStorage struct {
	gorm.DB
}

// NewSQLStorage is a convenience method to create a SQLStorage
func NewSQLStorage(dialect string, args ...interface{}) (*SQLStorage, error) {
	gormDB, err := gorm.Open(dialect, args...)
	if err != nil {
		return nil, err
	}
	return &SQLStorage{
		DB: gormDB,
	}, nil
}

// UpdateCurrent updates a single TUF.
func (db *SQLStorage) UpdateCurrent(gun string, update MetaUpdate) error {
	// ensure we're not inserting an immediately old version - can't use the
	// struct, because that only works with non-zero values, and Version
	// can be 0.
	exists := db.Where("gun = ? and role = ? and version >= ?",
		gun, update.Role, update.Version).First(&TUFFile{})

	if !exists.RecordNotFound() {
		return &ErrOldVersion{}
	}

	return db.Create(&TUFFile{
		Gun:     gun,
		Role:    update.Role,
		Version: update.Version,
		Data:    update.Data,
	}).Error
}

// UpdateMany atomically updates many TUF records in a single transaction
func (db *SQLStorage) UpdateMany(gun string, updates []MetaUpdate) error {
	tx := db.Begin()
	if tx.Error != nil {
		return tx.Error
	}

	rollback := func(err error) error {
		query := tx.Rollback()
		if query.Error != nil {
			logrus.Panic("Failed on Tx rollback with error: ", err.Error())
		}
		return err
	}

	var (
		query *gorm.DB
		added = make(map[uint]bool)
	)
	for _, update := range updates {
		// This looks like the same logic as UpdateCurrent, but if we just
		// called, version ordering in the updates list must be enforced
		// (you cannot insert the version 2 before version 1).  And we do
		// not care about monotonic ordering in the updates.
		query = db.Where("gun = ? and role = ? and version >= ?",
			gun, update.Role, update.Version).First(&TUFFile{})

		if !query.RecordNotFound() {
			return rollback(&ErrOldVersion{})
		}

		var row TUFFile
		query = tx.Where(map[string]interface{}{
			"gun":     gun,
			"role":    update.Role,
			"version": update.Version,
		}).Attrs("data", update.Data).FirstOrCreate(&row)

		if query.Error != nil {
			return rollback(query.Error)
		}
		// it's previously been added, which means it's a duplicate entry
		// in the same transaction
		if added[row.ID] {
			return rollback(&ErrOldVersion{})
		}
		added[row.ID] = true
	}
	return tx.Commit().Error
}

// GetCurrent gets a specific TUF record
func (db *SQLStorage) GetCurrent(gun, tufRole string) ([]byte, error) {
	var row TUFFile
	q := db.Select("data").Where(&TUFFile{Gun: gun, Role: tufRole}).Order("version desc").Limit(1).First(&row)

	if q.RecordNotFound() {
		return nil, &ErrNotFound{}
	} else if q.Error != nil {
		return nil, q.Error
	}
	return row.Data, nil
}

// Delete deletes all the records for a specific GUN
func (db *SQLStorage) Delete(gun string) error {
	return db.Where(&TUFFile{Gun: gun}).Delete(TUFFile{}).Error
}

// GetTimestampKey returns the timestamps Public Key data
func (db *SQLStorage) GetTimestampKey(gun string) (algorithm data.KeyAlgorithm, public []byte, err error) {
	logrus.Debug("retrieving timestamp key for ", gun)

	var row TimestampKey
	query := db.Select("cipher, public").Where(&TimestampKey{Gun: gun}).Find(&row)

	if query.RecordNotFound() {
		return "", nil, &ErrNoKey{gun: gun}
	} else if query.Error != nil {
		return "", nil, query.Error
	}

	return data.KeyAlgorithm(row.Cipher), row.Public, nil
}

// SetTimestampKey attempts to write a TimeStamp key and returns an error if it already exists
func (db *SQLStorage) SetTimestampKey(gun string, algorithm data.KeyAlgorithm, public []byte) error {

	entry := TimestampKey{
		Gun:    gun,
		Cipher: string(algorithm),
		Public: public,
	}

	if !db.Where(&entry).First(&TimestampKey{}).RecordNotFound() {
		return &ErrTimestampKeyExists{gun: gun}
	}

	return db.FirstOrCreate(&TimestampKey{}, &entry).Error
}
<commit_msg>Clean up SQLStorage docstring to refer to the model<commit_after>package storage

import (
	"github.com/Sirupsen/logrus"
	"github.com/endophage/gotuf/data"
	"github.com/jinzhu/gorm"
)

// SQLStorage implements a versioned store using a relational database.
// See server/storage/models.go
type SQLStorage struct {
	gorm.DB
}

// NewSQLStorage is a convenience method to create a SQLStorage
func NewSQLStorage(dialect string, args ...interface{}) (*SQLStorage, error) {
	gormDB, err := gorm.Open(dialect, args...)
	if err != nil {
		return nil, err
	}
	return &SQLStorage{
		DB: gormDB,
	}, nil
}

// UpdateCurrent updates a single TUF.
func (db *SQLStorage) UpdateCurrent(gun string, update MetaUpdate) error {
	// ensure we're not inserting an immediately old version - can't use the
	// struct, because that only works with non-zero values, and Version
	// can be 0.
	exists := db.Where("gun = ? and role = ? and version >= ?",
		gun, update.Role, update.Version).First(&TUFFile{})

	if !exists.RecordNotFound() {
		return &ErrOldVersion{}
	}

	return db.Create(&TUFFile{
		Gun:     gun,
		Role:    update.Role,
		Version: update.Version,
		Data:    update.Data,
	}).Error
}

// UpdateMany atomically updates many TUF records in a single transaction
func (db *SQLStorage) UpdateMany(gun string, updates []MetaUpdate) error {
	tx := db.Begin()
	if tx.Error != nil {
		return tx.Error
	}

	rollback := func(err error) error {
		query := tx.Rollback()
		if query.Error != nil {
			logrus.Panic("Failed on Tx rollback with error: ", err.Error())
		}
		return err
	}

	var (
		query *gorm.DB
		added = make(map[uint]bool)
	)
	for _, update := range updates {
		// This looks like the same logic as UpdateCurrent, but if we just
		// called, version ordering in the updates list must be enforced
		// (you cannot insert the version 2 before version 1).  And we do
		// not care about monotonic ordering in the updates.
		query = db.Where("gun = ? and role = ? and version >= ?",
			gun, update.Role, update.Version).First(&TUFFile{})

		if !query.RecordNotFound() {
			return rollback(&ErrOldVersion{})
		}

		var row TUFFile
		query = tx.Where(map[string]interface{}{
			"gun":     gun,
			"role":    update.Role,
			"version": update.Version,
		}).Attrs("data", update.Data).FirstOrCreate(&row)

		if query.Error != nil {
			return rollback(query.Error)
		}
		// it's previously been added, which means it's a duplicate entry
		// in the same transaction
		if added[row.ID] {
			return rollback(&ErrOldVersion{})
		}
		added[row.ID] = true
	}
	return tx.Commit().Error
}

// GetCurrent gets a specific TUF record
func (db *SQLStorage) GetCurrent(gun, tufRole string) ([]byte, error) {
	var row TUFFile
	q := db.Select("data").Where(&TUFFile{Gun: gun, Role: tufRole}).Order("version desc").Limit(1).First(&row)

	if q.RecordNotFound() {
		return nil, &ErrNotFound{}
	} else if q.Error != nil {
		return nil, q.Error
	}
	return row.Data, nil
}

// Delete deletes all the records for a specific GUN
func (db *SQLStorage) Delete(gun string) error {
	return db.Where(&TUFFile{Gun: gun}).Delete(TUFFile{}).Error
}

// GetTimestampKey returns the timestamps Public Key data
func (db *SQLStorage) GetTimestampKey(gun string) (algorithm data.KeyAlgorithm, public []byte, err error) {
	logrus.Debug("retrieving timestamp key for ", gun)

	var row TimestampKey
	query := db.Select("cipher, public").Where(&TimestampKey{Gun: gun}).Find(&row)

	if query.RecordNotFound() {
		return "", nil, &ErrNoKey{gun: gun}
	} else if query.Error != nil {
		return "", nil, query.Error
	}

	return data.KeyAlgorithm(row.Cipher), row.Public, nil
}

// SetTimestampKey attempts to write a TimeStamp key and returns an error if it already exists
func (db *SQLStorage) SetTimestampKey(gun string, algorithm data.KeyAlgorithm, public []byte) error {

	entry := TimestampKey{
		Gun:    gun,
		Cipher: string(algorithm),
		Public: public,
	}

	if !db.Where(&entry).First(&TimestampKey{}).RecordNotFound() {
		return &ErrTimestampKeyExists{gun: gun}
	}

	return db.FirstOrCreate(&TimestampKey{}, &entry).Error
}
<|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"
	"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.
*/

// NewBackupsOrigin returns a snapshot of where backup was run.
func NewBackupsOrigin(st *State, machine string) *metadata.Origin {
	// hostname could be derived from the environment...
	hostname, err := os.Hostname()
	if err != nil {
		panic(fmt.Sprintf("could not get hostname: %v", err))
	}
	origin := metadata.NewOrigin(
		st.EnvironTag().Id(),
		machine,
		hostname,
		version.Current.Number,
	)
	return origin
}

// 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) 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 "finished" fields.
	if doc.Finished == 0 {
		if doc.Checksum == "" {
			if doc.ChecksumFormat == "" {
				if doc.Size == 0 {
					if doc.Stored {
						return errors.New("Stored unexpected set")
					}
					return nil
				}
			}
		}
		return errors.New("missing finished")
	}
	if doc.Checksum == "" {
		return errors.New("missing finished")
	}
	if doc.ChecksumFormat == "" {
		return errors.New("missing finished")
	}
	if doc.Size == 0 {
		return errors.New("missing finished")
	}

	return nil
}

// asMetadata returns a new backups.Metadata based on the backupMetadataDoc.
func (doc *backupMetadataDoc) asMetadata() *metadata.Metadata {
	// Create a new Metadata.
	origin := metadata.NewOrigin(
		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)

	// Set the "finished" fields.
	if doc.Finished == 0 {
		if doc.Checksum == "" {
			if doc.ChecksumFormat == "" {
				if doc.Size == 0 {
					// Not finished so exit early.
					return meta
				}
			}
		}
	}
	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 {
		panic("invalid metadata doc")
	}
	if doc.Stored {
		meta.SetStored()
	}
	return meta
}

// updateFromMetadata copies the corresponding data from the metadata.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
}

// 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, err := utils.NewUUID()
	if err != nil {
		return "", errors.Annotate(err, "error generating new ID")
	}
	idStr := id.String()
	return idStr, addBackupMetadataID(st, metadata, idStr)
}

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

// 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) {
	metadata, ok := doc.(metadata.Metadata)
	if !ok {
		return "", errors.Errorf("doc must be of type state.metadata.Metadata")
	}
	return addBackupMetadata(s.state, &metadata)
}

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) {
	return nil, errors.Errorf("not implemented yet")
}

func (s *backupMetadataStorage) RemoveDoc(id string) error {
	return errors.Errorf("not implemented yet")
}

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{})

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>Add documentation for NewBackupsOrigin.<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"
	"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) 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 "finished" fields.
	if doc.Finished == 0 {
		if doc.Checksum == "" {
			if doc.ChecksumFormat == "" {
				if doc.Size == 0 {
					if doc.Stored {
						return errors.New("Stored unexpected set")
					}
					return nil
				}
			}
		}
		return errors.New("missing finished")
	}
	if doc.Checksum == "" {
		return errors.New("missing finished")
	}
	if doc.ChecksumFormat == "" {
		return errors.New("missing finished")
	}
	if doc.Size == 0 {
		return errors.New("missing finished")
	}

	return nil
}

// asMetadata returns a new backups.Metadata based on the backupMetadataDoc.
func (doc *backupMetadataDoc) asMetadata() *metadata.Metadata {
	// Create a new Metadata.
	origin := metadata.NewOrigin(
		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)

	// Set the "finished" fields.
	if doc.Finished == 0 {
		if doc.Checksum == "" {
			if doc.ChecksumFormat == "" {
				if doc.Size == 0 {
					// Not finished so exit early.
					return meta
				}
			}
		}
	}
	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 {
		panic("invalid metadata doc")
	}
	if doc.Stored {
		meta.SetStored()
	}
	return meta
}

// updateFromMetadata copies the corresponding data from the metadata.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
}

// 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, err := utils.NewUUID()
	if err != nil {
		return "", errors.Annotate(err, "error generating new ID")
	}
	idStr := id.String()
	return idStr, addBackupMetadataID(st, metadata, idStr)
}

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 {
		panic(fmt.Sprintf("could not get hostname: %v", err))
	}
	origin := metadata.NewOrigin(
		st.EnvironTag().Id(),
		machine,
		hostname,
		version.Current.Number,
	)
	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) {
	metadata, ok := doc.(metadata.Metadata)
	if !ok {
		return "", errors.Errorf("doc must be of type state.metadata.Metadata")
	}
	return addBackupMetadata(s.state, &metadata)
}

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) {
	return nil, errors.Errorf("not implemented yet")
}

func (s *backupMetadataStorage) RemoveDoc(id string) error {
	return errors.Errorf("not implemented yet")
}

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{})

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|>