<commit_before>//: ----------------------------------------------------------------------------
//: Copyright (C) 2017 Verizon.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    netflow_v9.go
//: details: netflow decoders handler
//: author:  Mehrdad Arshad Rad
//: date:    04/21/2017
//:
//: Licensed under the Apache License, Version 2.0 (the "License");
//: you may not use this file except in compliance with the License.
//: You may obtain a copy of the License at
//:
//:     http://www.apache.org/licenses/LICENSE-2.0
//:
//: Unless required by applicable law or agreed to in writing, software
//: distributed under the License is distributed on an "AS IS" BASIS,
//: WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//: See the License for the specific language governing permissions and
//: limitations under the License.
//: ----------------------------------------------------------------------------

package main

import (
	"bytes"
	"net"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/VerizonDigital/vflow/netflow/v9"
	"github.com/VerizonDigital/vflow/producer"
)

// NetflowV9 represents netflow v9 collector
type NetflowV9 struct {
	port    int
	addr    string
	workers int
	stop    bool
	stats   NetflowV9Stats
	pool    chan chan struct{}
}

// NetflowV9UDPMsg represents netflow v9 UDP data
type NetflowV9UDPMsg struct {
	raddr *net.UDPAddr
	body  []byte
}

// NetflowV9Stats represents netflow v9 stats
type NetflowV9Stats struct {
	UDPQueue     int
	MessageQueue int
	UDPCount     uint64
	DecodedCount uint64
	MQErrorCount uint64
	Workers      int32
}

var (
	netflowV9UDPCh = make(chan NetflowV9UDPMsg, 1000)
	netflowV9MQCh  = make(chan []byte, 1000)

	mCacheNF9 netflow9.MemCache

	// ipfix udp payload pool
	netflowV9Buffer = &sync.Pool{
		New: func() interface{} {
			return make([]byte, opts.NetflowV9UDPSize)
		},
	}
)

// NewNetflowV9 constructs NetflowV9
func NewNetflowV9() *NetflowV9 {
	return &NetflowV9{
		port:    opts.IPFIXPort,
		workers: opts.IPFIXWorkers,
		pool:    make(chan chan struct{}, maxWorkers),
	}
}

func (i *NetflowV9) run() {
	//TODO
	// exit if the ipfix is disabled
	if !opts.NetflowV9Enabled {
		logger.Println("netflowv9 has been disabled")
		return
	}

	hostPort := net.JoinHostPort(i.addr, strconv.Itoa(i.port))
	udpAddr, _ := net.ResolveUDPAddr("udp", hostPort)

	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		logger.Fatal(err)
	}

	atomic.AddInt32(&i.stats.Workers, int32(i.workers))
	for n := 0; n < i.workers; n++ {
		go func() {
			wQuit := make(chan struct{})
			i.pool <- wQuit
			i.netflowV9Worker(wQuit)
		}()
	}

	logger.Printf("netflow v9 is running (workers#: %d)", i.workers)

	go func() {
		p := producer.NewProducer(opts.MQName)

		p.MQConfigFile = opts.MQConfigFile
		p.MQErrorCount = &i.stats.MQErrorCount
		p.Logger = logger
		p.Chan = netflowV9MQCh
		p.Topic = opts.NetflowV9Topic

		if err := p.Run(); err != nil {
			logger.Fatal(err)
		}
	}()

	go func() {
		if !opts.DynWorkers {
			logger.Println("netflow v9 dynamic worker disabled")
			return
		}

		i.dynWorkers()
	}()

	for !i.stop {
		b := netflowV9Buffer.Get().([]byte)
		conn.SetReadDeadline(time.Now().Add(1e9))
		n, raddr, err := conn.ReadFromUDP(b)
		if err != nil {
			continue
		}
		atomic.AddUint64(&i.stats.UDPCount, 1)
		netflowV9UDPCh <- NetflowV9UDPMsg{raddr, b[:n]}
	}

}

func (i *NetflowV9) shutdown() {
	// exit if the netflow v9 is disabled
	if !opts.NetflowV9Enabled {
		logger.Println("netflow v9 disabled")
		return
	}

	// stop reading from UDP listener
	i.stop = true
	logger.Println("stopping netflow v9 service gracefully ...")
	time.Sleep(1 * time.Second)

	// dump the templates to storage
	if err := mCacheNF9.Dump(opts.NetflowV9TplCacheFile); err != nil {
		logger.Println("couldn't not dump template", err)
	}

	// logging and close UDP channel
	logger.Println("netflow v9 has been shutdown")
	close(netflowV9UDPCh)
}

func (i *NetflowV9) netflowV9Worker(wQuit chan struct{}) {
	var (
		decodedMsg *netflow9.Message
		msg        = NetflowV9UDPMsg{body: netflowV9Buffer.Get().([]byte)}
		buf        = new(bytes.Buffer)
		err        error
		ok         bool
		b          []byte
	)

LOOP:
	for {

		netflowV9Buffer.Put(msg.body[:opts.NetflowV9UDPSize])
		buf.Reset()

		select {
		case <-wQuit:
			break LOOP
		case msg, ok = <-netflowV9UDPCh:
			if !ok {
				break LOOP
			}
		}

		if opts.Verbose {
			logger.Printf("rcvd netflow v9 data from: %s, size: %d bytes",
				msg.raddr, len(msg.body))
		}

		d := netflow9.NewDecoder(msg.raddr.IP, msg.body)
		if decodedMsg, err = d.Decode(mCacheNF9); err != nil {
			logger.Println(err)
			continue
		}

		atomic.AddUint64(&i.stats.DecodedCount, 1)

		if decodedMsg.DataSets != nil {
			b, err = decodedMsg.JSONMarshal(buf)
			if err != nil {
				logger.Println(err)
				continue
			}

			select {
			case netflowV9MQCh <- b:
			default:
			}
		}

		if opts.Verbose {
			logger.Println(string(b))
		}

	}

}

func (i *NetflowV9) status() *NetflowV9Stats {
	return &NetflowV9Stats{
		UDPQueue:     len(netflowV9UDPCh),
		MessageQueue: len(netflowV9MQCh),
		UDPCount:     atomic.LoadUint64(&i.stats.UDPCount),
		DecodedCount: atomic.LoadUint64(&i.stats.DecodedCount),
		MQErrorCount: atomic.LoadUint64(&i.stats.MQErrorCount),
		Workers:      atomic.LoadInt32(&i.stats.Workers),
	}

}

func (i *NetflowV9) dynWorkers() {
	var load, nSeq, newWorkers, workers, n int

	tick := time.Tick(120 * time.Second)

	for {
		<-tick
		load = 0

		for n = 0; n < 30; n++ {
			time.Sleep(1 * time.Second)
			load += len(netflowV9UDPCh)
		}

		if load > 15 {

			switch {
			case load > 300:
				newWorkers = 100
			case load > 200:
				newWorkers = 60
			case load > 100:
				newWorkers = 40
			default:
				newWorkers = 30
			}

			workers = int(atomic.LoadInt32(&i.stats.Workers))
			if workers+newWorkers > maxWorkers {
				logger.Println("netflow v9 :: max out workers")
				continue
			}

			for n = 0; n < newWorkers; n++ {
				go func() {
					atomic.AddInt32(&i.stats.Workers, 1)
					wQuit := make(chan struct{})
					i.pool <- wQuit
					i.netflowV9Worker(wQuit)
				}()
			}

		}

		if load == 0 {
			nSeq++
		} else {
			nSeq = 0
			continue
		}

		if nSeq > 15 {
			for n = 0; n < 10; n++ {
				if len(i.pool) > i.workers {
					atomic.AddInt32(&i.stats.Workers, -1)
					wQuit := <-i.pool
					close(wQuit)
				}
			}

			nSeq = 0
		}
	}

}
<commit_msg> fix new netflow object<commit_after>//: ----------------------------------------------------------------------------
//: Copyright (C) 2017 Verizon.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    netflow_v9.go
//: details: netflow decoders handler
//: author:  Mehrdad Arshad Rad
//: date:    04/21/2017
//:
//: Licensed under the Apache License, Version 2.0 (the "License");
//: you may not use this file except in compliance with the License.
//: You may obtain a copy of the License at
//:
//:     http://www.apache.org/licenses/LICENSE-2.0
//:
//: Unless required by applicable law or agreed to in writing, software
//: distributed under the License is distributed on an "AS IS" BASIS,
//: WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//: See the License for the specific language governing permissions and
//: limitations under the License.
//: ----------------------------------------------------------------------------

package main

import (
	"bytes"
	"net"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/VerizonDigital/vflow/netflow/v9"
	"github.com/VerizonDigital/vflow/producer"
)

// NetflowV9 represents netflow v9 collector
type NetflowV9 struct {
	port    int
	addr    string
	workers int
	stop    bool
	stats   NetflowV9Stats
	pool    chan chan struct{}
}

// NetflowV9UDPMsg represents netflow v9 UDP data
type NetflowV9UDPMsg struct {
	raddr *net.UDPAddr
	body  []byte
}

// NetflowV9Stats represents netflow v9 stats
type NetflowV9Stats struct {
	UDPQueue     int
	MessageQueue int
	UDPCount     uint64
	DecodedCount uint64
	MQErrorCount uint64
	Workers      int32
}

var (
	netflowV9UDPCh = make(chan NetflowV9UDPMsg, 1000)
	netflowV9MQCh  = make(chan []byte, 1000)

	mCacheNF9 netflow9.MemCache

	// ipfix udp payload pool
	netflowV9Buffer = &sync.Pool{
		New: func() interface{} {
			return make([]byte, opts.NetflowV9UDPSize)
		},
	}
)

// NewNetflowV9 constructs NetflowV9
func NewNetflowV9() *NetflowV9 {
	return &NetflowV9{
		port:    opts.NetflowV9Port,
		workers: opts.NetflowV9Workers,
		pool:    make(chan chan struct{}, maxWorkers),
	}
}

func (i *NetflowV9) run() {
	//TODO
	// exit if the ipfix is disabled
	if !opts.NetflowV9Enabled {
		logger.Println("netflowv9 has been disabled")
		return
	}

	hostPort := net.JoinHostPort(i.addr, strconv.Itoa(i.port))
	udpAddr, _ := net.ResolveUDPAddr("udp", hostPort)

	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		logger.Fatal(err)
	}

	atomic.AddInt32(&i.stats.Workers, int32(i.workers))
	for n := 0; n < i.workers; n++ {
		go func() {
			wQuit := make(chan struct{})
			i.pool <- wQuit
			i.netflowV9Worker(wQuit)
		}()
	}

	logger.Printf("netflow v9 is running (workers#: %d)", i.workers)

	go func() {
		p := producer.NewProducer(opts.MQName)

		p.MQConfigFile = opts.MQConfigFile
		p.MQErrorCount = &i.stats.MQErrorCount
		p.Logger = logger
		p.Chan = netflowV9MQCh
		p.Topic = opts.NetflowV9Topic

		if err := p.Run(); err != nil {
			logger.Fatal(err)
		}
	}()

	go func() {
		if !opts.DynWorkers {
			logger.Println("netflow v9 dynamic worker disabled")
			return
		}

		i.dynWorkers()
	}()

	for !i.stop {
		b := netflowV9Buffer.Get().([]byte)
		conn.SetReadDeadline(time.Now().Add(1e9))
		n, raddr, err := conn.ReadFromUDP(b)
		if err != nil {
			continue
		}
		atomic.AddUint64(&i.stats.UDPCount, 1)
		netflowV9UDPCh <- NetflowV9UDPMsg{raddr, b[:n]}
	}

}

func (i *NetflowV9) shutdown() {
	// exit if the netflow v9 is disabled
	if !opts.NetflowV9Enabled {
		logger.Println("netflow v9 disabled")
		return
	}

	// stop reading from UDP listener
	i.stop = true
	logger.Println("stopping netflow v9 service gracefully ...")
	time.Sleep(1 * time.Second)

	// dump the templates to storage
	if err := mCacheNF9.Dump(opts.NetflowV9TplCacheFile); err != nil {
		logger.Println("couldn't not dump template", err)
	}

	// logging and close UDP channel
	logger.Println("netflow v9 has been shutdown")
	close(netflowV9UDPCh)
}

func (i *NetflowV9) netflowV9Worker(wQuit chan struct{}) {
	var (
		decodedMsg *netflow9.Message
		msg        = NetflowV9UDPMsg{body: netflowV9Buffer.Get().([]byte)}
		buf        = new(bytes.Buffer)
		err        error
		ok         bool
		b          []byte
	)

LOOP:
	for {

		netflowV9Buffer.Put(msg.body[:opts.NetflowV9UDPSize])
		buf.Reset()

		select {
		case <-wQuit:
			break LOOP
		case msg, ok = <-netflowV9UDPCh:
			if !ok {
				break LOOP
			}
		}

		if opts.Verbose {
			logger.Printf("rcvd netflow v9 data from: %s, size: %d bytes",
				msg.raddr, len(msg.body))
		}

		d := netflow9.NewDecoder(msg.raddr.IP, msg.body)
		if decodedMsg, err = d.Decode(mCacheNF9); err != nil {
			logger.Println(err)
			continue
		}

		atomic.AddUint64(&i.stats.DecodedCount, 1)

		if decodedMsg.DataSets != nil {
			b, err = decodedMsg.JSONMarshal(buf)
			if err != nil {
				logger.Println(err)
				continue
			}

			select {
			case netflowV9MQCh <- b:
			default:
			}
		}

		if opts.Verbose {
			logger.Println(string(b))
		}

	}

}

func (i *NetflowV9) status() *NetflowV9Stats {
	return &NetflowV9Stats{
		UDPQueue:     len(netflowV9UDPCh),
		MessageQueue: len(netflowV9MQCh),
		UDPCount:     atomic.LoadUint64(&i.stats.UDPCount),
		DecodedCount: atomic.LoadUint64(&i.stats.DecodedCount),
		MQErrorCount: atomic.LoadUint64(&i.stats.MQErrorCount),
		Workers:      atomic.LoadInt32(&i.stats.Workers),
	}

}

func (i *NetflowV9) dynWorkers() {
	var load, nSeq, newWorkers, workers, n int

	tick := time.Tick(120 * time.Second)

	for {
		<-tick
		load = 0

		for n = 0; n < 30; n++ {
			time.Sleep(1 * time.Second)
			load += len(netflowV9UDPCh)
		}

		if load > 15 {

			switch {
			case load > 300:
				newWorkers = 100
			case load > 200:
				newWorkers = 60
			case load > 100:
				newWorkers = 40
			default:
				newWorkers = 30
			}

			workers = int(atomic.LoadInt32(&i.stats.Workers))
			if workers+newWorkers > maxWorkers {
				logger.Println("netflow v9 :: max out workers")
				continue
			}

			for n = 0; n < newWorkers; n++ {
				go func() {
					atomic.AddInt32(&i.stats.Workers, 1)
					wQuit := make(chan struct{})
					i.pool <- wQuit
					i.netflowV9Worker(wQuit)
				}()
			}

		}

		if load == 0 {
			nSeq++
		} else {
			nSeq = 0
			continue
		}

		if nSeq > 15 {
			for n = 0; n < 10; n++ {
				if len(i.pool) > i.workers {
					atomic.AddInt32(&i.stats.Workers, -1)
					wQuit := <-i.pool
					close(wQuit)
				}
			}

			nSeq = 0
		}
	}

}
<|endoftext|>


<commit_before>package video

import "image"

var color_channel = [12]int{0, 0, 0, 0, 1, 2, 1, 2, 1, 2, 1, 2}

type Macroblock struct {
	macroblock_address_increment uint32
	macroblock_type              *MacroblockType
	spatial_temporal_weight_code uint32
	frame_motion_type            uint32
	field_motion_type            uint32
	dct_type                     bool
	quantiser_scale_code         uint32

	cpb int
}

func (br *VideoSequence) macroblock(mbAddress int, frameSlice *image.YCbCr) (int, error) {

	mb := Macroblock{}

	for {
		if nextbits, err := br.Peek32(11); err != nil {
			return 0, err
		} else if nextbits == 0x08 { // 0000 0001 000
			br.Trash(11)
			mb.macroblock_address_increment += 33
		} else {
			break
		}
	}

	if incr, err := macroblockAddressIncrementDecoder.Decode(br); err != nil {
		return 0, err
	} else {
		mb.macroblock_address_increment += incr
	}

	if mb.macroblock_address_increment > 1 {
		br.resetDCPredictors()
	}

	// Reset motion vector predictors: P-picture with a skipped macroblock (7.6.4.3)
	if br.PictureHeader.picture_coding_type == PFrame &&
		mb.macroblock_address_increment > 1 {
		br.pMV.reset()
	}

	mbAddress += int(mb.macroblock_address_increment)

	if err := br.macroblock_mode(&mb); err != nil {
		return 0, err
	}

	if mb.macroblock_type.macroblock_intra == false {
		br.resetDCPredictors()
	}

	// Reset motion vector predictors: intra macroblock without concealment motion vectors (7.6.4.3)
	if mb.macroblock_type.macroblock_intra == true &&
		br.PictureCodingExtension.concealment_motion_vectors {
		br.pMV.reset()
	}

	// Reset motion vector predictors: non-intra P-picture with no forward motion vectors (7.6.4.3)
	if br.PictureHeader.picture_coding_type == PFrame &&
		mb.macroblock_type.macroblock_intra == false &&
		mb.macroblock_type.macroblock_motion_forward == false {
		br.pMV.reset()
	}

	if mb.macroblock_type.macroblock_quant {
		if qsc, err := br.Read32(5); err != nil {
			return 0, err
		} else {
			mb.quantiser_scale_code = qsc
			br.currentQSC = qsc
		}
	}

	var mvd motionVectorData

	if mb.macroblock_type.macroblock_motion_forward ||
		(mb.macroblock_type.macroblock_intra && br.PictureCodingExtension.concealment_motion_vectors) {
		if err := br.motion_vectors(0, &mb, &mvd); err != nil {
			return 0, err
		}
	}

	if mb.macroblock_type.macroblock_motion_backward {
		if err := br.motion_vectors(1, &mb, &mvd); err != nil {
			return 0, err
		}
	}

	if mb.macroblock_type.macroblock_intra && br.PictureCodingExtension.concealment_motion_vectors {
		if err := marker_bit(br); err != nil {
			return 0, err
		}
	}

	if mb.macroblock_type.macroblock_pattern {
		if cpb, err := coded_block_pattern(br, br.SequenceExtension.chroma_format); err != nil {
			return 0, nil
		} else {
			mb.cpb = cpb
		}
	}

	var block_count int
	switch br.SequenceExtension.chroma_format {
	case ChromaFormat_420:
		block_count = 6
	case ChromaFormat_422:
		block_count = 8
	case ChromaFormat_444:
		block_count = 12
	}

	pattern_code := mb.decodePatternCode(br.SequenceExtension.chroma_format)

	var b block
	var decoded block

	for i := 0; i < block_count; i++ {
		cc := color_channel[i]

		if pattern_code[i] {
			if err := br.block(cc, &mb, &b); err != nil {
				return 0, err
			}
		}

		err := br.decode_block(cc, mb.macroblock_type.macroblock_intra, &b, &decoded)
		if err != nil {
			return 0, err
		}
		idct(&decoded)
		updateFrameSlice(i, mbAddress, mb.dct_type, frameSlice, &decoded)
	}

	return mbAddress, nil
}

type clampedBlock [blockSize]uint8

func clamp(dest *clampedBlock, src *block) {
	for i := 0; i < 64; i++ {
		if src[i] > 255 {
			dest[i] = 255
		} else if src[i] < 0 {
			dest[i] = 0
		} else {
			dest[i] = uint8(src[i])
		}
	}
}

func updateFrameSlice(i, mbAddress int, interlaced bool, frameSlice *image.YCbCr, b *block) {

	var cb clampedBlock
	clamp(&cb, b)

	if interlaced {
		switch i {
		case 0:
			xs := (mbAddress - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := (y*2)*frameSlice.YStride + xs
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 1:
			xs := (mbAddress - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := (y*2)*frameSlice.YStride + (xs + 8)
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 2:
			xs := (mbAddress - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := ((y*2)+1)*frameSlice.YStride + xs
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 3:
			xs := (mbAddress - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := ((y*2)+1)*frameSlice.YStride + (xs + 8)
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 4:
			xs := (mbAddress - 1) * 8
			for y := 0; y < 8; y++ {
				si := y * 8
				di := y*frameSlice.CStride + xs
				copy(frameSlice.Cb[di:di+8], cb[si:si+8])
			}
		case 5:
			xs := (mbAddress - 1) * 8
			for y := 0; y < 8; y++ {
				si := y * 8
				di := y*frameSlice.CStride + xs
				copy(frameSlice.Cr[di:di+8], cb[si:si+8])
			}
		}
	} else {
		switch i {
		case 0:
			xs := (mbAddress - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := y*frameSlice.YStride + xs
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 1:
			xs := (mbAddress - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := y*frameSlice.YStride + (xs + 8)
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 2:
			xs := (mbAddress - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := (y+8)*frameSlice.YStride + xs
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 3:
			xs := (mbAddress - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := (y+8)*frameSlice.YStride + (xs + 8)
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 4:
			xs := (mbAddress - 1) * 8
			for y := 0; y < 8; y++ {
				si := y * 8
				di := y*frameSlice.CStride + xs
				copy(frameSlice.Cb[di:di+8], cb[si:si+8])
			}
		case 5:
			xs := (mbAddress - 1) * 8
			for y := 0; y < 8; y++ {
				si := y * 8
				di := y*frameSlice.CStride + xs
				copy(frameSlice.Cr[di:di+8], cb[si:si+8])
			}
		}
	}

}

type PatternCode [12]bool

func (mb *Macroblock) decodePatternCode(chroma_format chromaFormat) (pattern_code PatternCode) {
	for i := 0; i < 12; i++ {
		if mb.macroblock_type.macroblock_intra {
			pattern_code[i] = true
		} else {
			pattern_code[i] = false
		}
	}

	if mb.macroblock_type.macroblock_pattern {
		for i := 0; i < 6; i++ {
			mask := 1 << uint(5-i)
			if mb.cpb&mask == mask {
				pattern_code[i] = true
			}
		}

		if chroma_format == ChromaFormat_422 || chroma_format == ChromaFormat_444 {
			panic("unsupported: coded block pattern chroma format")
		}
	}

	return
}

func (br *VideoSequence) macroblock_mode(mb *Macroblock) (err error) {

	var typeDecoder macroblockTypeDecoderFn
	switch br.PictureHeader.picture_coding_type {
	case IntraCoded:
		typeDecoder = macroblockTypeDecoder.IFrame
	case PredictiveCoded:
		typeDecoder = macroblockTypeDecoder.PFrame
	case BidirectionallyPredictiveCoded:
		typeDecoder = macroblockTypeDecoder.BFrame
	default:
		panic("not implemented: macroblock type decoder")
	}

	mb.macroblock_type, err = typeDecoder(br)
	if err != nil {
		return err
	}

	if mb.macroblock_type.spatial_temporal_weight_code_flag &&
		false /* ( spatial_temporal_weight_code_table_index != ‘00’) */ {
		mb.spatial_temporal_weight_code, err = br.Read32(2)
		if err != nil {
			return err
		}
	}

	if mb.macroblock_type.macroblock_motion_forward ||
		mb.macroblock_type.macroblock_motion_backward {
		if br.PictureCodingExtension.picture_structure == PictureStructure_FramePicture {
			if br.PictureCodingExtension.frame_pred_frame_dct == 0 {
				mb.frame_motion_type, err = br.Read32(2)
				if err != nil {
					return err
				}
			}
		} else {
			mb.field_motion_type, err = br.Read32(2)
			if err != nil {
				return err
			}
		}
	}

	if br.PictureCodingExtension.picture_structure == PictureStructure_FramePicture &&
		br.PictureCodingExtension.frame_pred_frame_dct == 0 &&
		(mb.macroblock_type.macroblock_intra || mb.macroblock_type.macroblock_pattern) {
		mb.dct_type, err = br.ReadBit() //dct_type 1 uimsbf
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>renaming mb_address<commit_after>package video

import "image"

var color_channel = [12]int{0, 0, 0, 0, 1, 2, 1, 2, 1, 2, 1, 2}

type Macroblock struct {
	macroblock_address_increment uint32
	macroblock_type              *MacroblockType
	spatial_temporal_weight_code uint32
	frame_motion_type            uint32
	field_motion_type            uint32
	dct_type                     bool
	quantiser_scale_code         uint32

	cpb int
}

func (br *VideoSequence) macroblock(mb_row int, mb_address int, frameSlice *image.YCbCr) (int, error) {

	mb := Macroblock{}

	for {
		if nextbits, err := br.Peek32(11); err != nil {
			return 0, err
		} else if nextbits == 0x08 { // 0000 0001 000
			br.Trash(11)
			mb.macroblock_address_increment += 33
		} else {
			break
		}
	}

	if incr, err := macroblockAddressIncrementDecoder.Decode(br); err != nil {
		return 0, err
	} else {
		mb.macroblock_address_increment += incr
	}

	if mb.macroblock_address_increment > 1 {
		br.resetDCPredictors()
	}

	// Reset motion vector predictors: P-picture with a skipped macroblock (7.6.4.3)
	if br.PictureHeader.picture_coding_type == PFrame &&
		mb.macroblock_address_increment > 1 {
		br.pMV.reset()
	}

	mb_address += int(mb.macroblock_address_increment)

	if err := br.macroblock_mode(&mb); err != nil {
		return 0, err
	}

	if mb.macroblock_type.macroblock_intra == false {
		br.resetDCPredictors()
	}

	// Reset motion vector predictors: intra macroblock without concealment motion vectors (7.6.4.3)
	if mb.macroblock_type.macroblock_intra == true &&
		br.PictureCodingExtension.concealment_motion_vectors {
		br.pMV.reset()
	}

	// Reset motion vector predictors: non-intra P-picture with no forward motion vectors (7.6.4.3)
	if br.PictureHeader.picture_coding_type == PFrame &&
		mb.macroblock_type.macroblock_intra == false &&
		mb.macroblock_type.macroblock_motion_forward == false {
		br.pMV.reset()
	}

	if mb.macroblock_type.macroblock_quant {
		if qsc, err := br.Read32(5); err != nil {
			return 0, err
		} else {
			mb.quantiser_scale_code = qsc
			br.currentQSC = qsc
		}
	}

	var mvd motionVectorData

	if mb.macroblock_type.macroblock_motion_forward ||
		(mb.macroblock_type.macroblock_intra && br.PictureCodingExtension.concealment_motion_vectors) {
		if err := br.motion_vectors(0, &mb, &mvd); err != nil {
			return 0, err
		}
	}

	if mb.macroblock_type.macroblock_motion_backward {
		if err := br.motion_vectors(1, &mb, &mvd); err != nil {
			return 0, err
		}
	}

	if mb.macroblock_type.macroblock_intra && br.PictureCodingExtension.concealment_motion_vectors {
		if err := marker_bit(br); err != nil {
			return 0, err
		}
	}

	if mb.macroblock_type.macroblock_pattern {
		if cpb, err := coded_block_pattern(br, br.SequenceExtension.chroma_format); err != nil {
			return 0, nil
		} else {
			mb.cpb = cpb
		}
	}

	var block_count int
	switch br.SequenceExtension.chroma_format {
	case ChromaFormat_420:
		block_count = 6
	case ChromaFormat_422:
		block_count = 8
	case ChromaFormat_444:
		block_count = 12
	}

	pattern_code := mb.decodePatternCode(br.SequenceExtension.chroma_format)

	var b block
	var decoded block

	for i := 0; i < block_count; i++ {
		cc := color_channel[i]

		if pattern_code[i] {
			if err := br.block(cc, &mb, &b); err != nil {
				return 0, err
			}
		}

		err := br.decode_block(cc, mb.macroblock_type.macroblock_intra, &b, &decoded)
		if err != nil {
			return 0, err
		}
		idct(&decoded)
		updateFrameSlice(i, mb_address, mb.dct_type, frameSlice, &decoded)
	}

	return mb_address, nil
}

type clampedBlock [blockSize]uint8

func clamp(dest *clampedBlock, src *block) {
	for i := 0; i < 64; i++ {
		if src[i] > 255 {
			dest[i] = 255
		} else if src[i] < 0 {
			dest[i] = 0
		} else {
			dest[i] = uint8(src[i])
		}
	}
}

func updateFrameSlice(i, mb_address int, interlaced bool, frameSlice *image.YCbCr, b *block) {

	var cb clampedBlock
	clamp(&cb, b)

	if interlaced {
		switch i {
		case 0:
			xs := (mb_address - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := (y*2)*frameSlice.YStride + xs
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 1:
			xs := (mb_address - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := (y*2)*frameSlice.YStride + (xs + 8)
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 2:
			xs := (mb_address - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := ((y*2)+1)*frameSlice.YStride + xs
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 3:
			xs := (mb_address - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := ((y*2)+1)*frameSlice.YStride + (xs + 8)
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 4:
			xs := (mb_address - 1) * 8
			for y := 0; y < 8; y++ {
				si := y * 8
				di := y*frameSlice.CStride + xs
				copy(frameSlice.Cb[di:di+8], cb[si:si+8])
			}
		case 5:
			xs := (mb_address - 1) * 8
			for y := 0; y < 8; y++ {
				si := y * 8
				di := y*frameSlice.CStride + xs
				copy(frameSlice.Cr[di:di+8], cb[si:si+8])
			}
		}
	} else {
		switch i {
		case 0:
			xs := (mb_address - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := y*frameSlice.YStride + xs
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 1:
			xs := (mb_address - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := y*frameSlice.YStride + (xs + 8)
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 2:
			xs := (mb_address - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := (y+8)*frameSlice.YStride + xs
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 3:
			xs := (mb_address - 1) * 16
			for y := 0; y < 8; y++ {
				si := y * 8
				di := (y+8)*frameSlice.YStride + (xs + 8)
				copy(frameSlice.Y[di:di+8], cb[si:si+8])
			}
		case 4:
			xs := (mb_address - 1) * 8
			for y := 0; y < 8; y++ {
				si := y * 8
				di := y*frameSlice.CStride + xs
				copy(frameSlice.Cb[di:di+8], cb[si:si+8])
			}
		case 5:
			xs := (mb_address - 1) * 8
			for y := 0; y < 8; y++ {
				si := y * 8
				di := y*frameSlice.CStride + xs
				copy(frameSlice.Cr[di:di+8], cb[si:si+8])
			}
		}
	}

}

type PatternCode [12]bool

func (mb *Macroblock) decodePatternCode(chroma_format chromaFormat) (pattern_code PatternCode) {
	for i := 0; i < 12; i++ {
		if mb.macroblock_type.macroblock_intra {
			pattern_code[i] = true
		} else {
			pattern_code[i] = false
		}
	}

	if mb.macroblock_type.macroblock_pattern {
		for i := 0; i < 6; i++ {
			mask := 1 << uint(5-i)
			if mb.cpb&mask == mask {
				pattern_code[i] = true
			}
		}

		if chroma_format == ChromaFormat_422 || chroma_format == ChromaFormat_444 {
			panic("unsupported: coded block pattern chroma format")
		}
	}

	return
}

func (br *VideoSequence) macroblock_mode(mb *Macroblock) (err error) {

	var typeDecoder macroblockTypeDecoderFn
	switch br.PictureHeader.picture_coding_type {
	case IntraCoded:
		typeDecoder = macroblockTypeDecoder.IFrame
	case PredictiveCoded:
		typeDecoder = macroblockTypeDecoder.PFrame
	case BidirectionallyPredictiveCoded:
		typeDecoder = macroblockTypeDecoder.BFrame
	default:
		panic("not implemented: macroblock type decoder")
	}

	mb.macroblock_type, err = typeDecoder(br)
	if err != nil {
		return err
	}

	if mb.macroblock_type.spatial_temporal_weight_code_flag &&
		false /* ( spatial_temporal_weight_code_table_index != ‘00’) */ {
		mb.spatial_temporal_weight_code, err = br.Read32(2)
		if err != nil {
			return err
		}
	}

	if mb.macroblock_type.macroblock_motion_forward ||
		mb.macroblock_type.macroblock_motion_backward {
		if br.PictureCodingExtension.picture_structure == PictureStructure_FramePicture {
			if br.PictureCodingExtension.frame_pred_frame_dct == 0 {
				mb.frame_motion_type, err = br.Read32(2)
				if err != nil {
					return err
				}
			}
		} else {
			mb.field_motion_type, err = br.Read32(2)
			if err != nil {
				return err
			}
		}
	}

	if br.PictureCodingExtension.picture_structure == PictureStructure_FramePicture &&
		br.PictureCodingExtension.frame_pred_frame_dct == 0 &&
		(mb.macroblock_type.macroblock_intra || mb.macroblock_type.macroblock_pattern) {
		mb.dct_type, err = br.ReadBit() //dct_type 1 uimsbf
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package ssh

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/pkg/errors"

	"github.com/cockroachdb/roachprod/config"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/crypto/ssh/knownhosts"
)

var knownHosts ssh.HostKeyCallback
var knownHostsOnce sync.Once
var InsecureIgnoreHostKey bool

func getKnownHosts() ssh.HostKeyCallback {
	knownHostsOnce.Do(func() {
		var err error
		if InsecureIgnoreHostKey {
			knownHosts = ssh.InsecureIgnoreHostKey()
		} else {
			knownHosts, err = knownhosts.New(filepath.Join(os.Getenv("HOME"), ".ssh", "known_hosts"))
			if err != nil {
				log.Fatal(err)
			}
		}
	})
	return knownHosts
}

func getSSHAgentSigners() []ssh.Signer {
	const authSockEnv = "SSH_AUTH_SOCK"
	agentSocket := os.Getenv(authSockEnv)
	if agentSocket == "" {
		return nil
	}
	sock, err := net.Dial("unix", agentSocket)
	if err != nil {
		log.Printf("SSH_AUTH_SOCK set but unable to connect to agent: %s", err)
		return nil
	}
	agent := agent.NewClient(sock)
	signers, err := agent.Signers()
	if err != nil {
		log.Printf("unable to retrieve keys from agent: %s", err)
		return nil
	}
	return signers
}

func getSSHKeySigner(path string) ssh.Signer {
	key, err := ioutil.ReadFile(path)
	if err != nil {
		if !os.IsNotExist(err) {
			log.Printf("unable to read SSH key %q: %s", path, err)
		}
		return nil
	}

	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		if strings.Contains(err.Error(), "cannot decode encrypted private key") {
			log.Printf(
				"skipping encrypted SSH key %q; if necessary, add the key to your SSH agent", path)
		} else {
			log.Printf("unable to parse SSH key %q: %s", path, err)
		}
		return nil
	}
	return signer
}

func getDefaultSSHKeySigners() []ssh.Signer {
	var signers []ssh.Signer
	for _, name := range []string{"id_rsa", "google_compute_engine"} {
		s := getSSHKeySigner(filepath.Join(config.OSUser.HomeDir, ".ssh", name))
		if s != nil {
			signers = append(signers, s)
		}
	}
	return signers
}

func newSSHClient(user, host string) (*ssh.Client, net.Conn, error) {
	config := &ssh.ClientConfig{
		User:            user,
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(sshState.signers...)},
		HostKeyCallback: getKnownHosts(),
	}
	config.SetDefaults()

	addr := fmt.Sprintf("%s:22", host)
	conn, err := net.DialTimeout("tcp", addr, 30*time.Second)
	if err != nil {
		return nil, nil, err
	}
	c, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
	if err != nil {
		return nil, nil, err
	}
	return ssh.NewClient(c, chans, reqs), conn, nil
}

type sshClient struct {
	sync.Mutex
	*ssh.Client
}

var sshState = struct {
	signers     []ssh.Signer
	signersInit sync.Once

	clients  map[string]*sshClient
	clientMu sync.Mutex
}{
	clients: map[string]*sshClient{},
}

func NewSSHSession(user, host string) (*ssh.Session, error) {
	sshState.clientMu.Lock()
	target := fmt.Sprintf("%s@%s", user, host)
	client := sshState.clients[target]
	if client == nil {
		client = &sshClient{}
		sshState.clients[target] = client
	}
	sshState.clientMu.Unlock()

	sshState.signersInit.Do(func() {
		sshState.signers = append(sshState.signers, getSSHAgentSigners()...)
		sshState.signers = append(sshState.signers, getDefaultSSHKeySigners()...)
	})

	client.Lock()
	defer client.Unlock()
	if client.Client == nil {
		var err error
		client.Client, _, err = newSSHClient(user, host)
		if err != nil {
			return nil, err
		}
	}
	return client.NewSession()
}

func IsSigKill(err error) bool {
	switch t := err.(type) {
	case *ssh.ExitError:
		return t.Signal() == string(ssh.SIGKILL)
	}
	return false
}

type progressWriter struct {
	writer   io.Writer
	done     int64
	total    int64
	progress func(float64)
}

func (p *progressWriter) Write(b []byte) (int, error) {
	n, err := p.writer.Write(b)
	if err == nil {
		p.done += int64(n)
		p.progress(float64(p.done) / float64(p.total))
	}
	return n, err
}

func SCPPut(src, dest string, progress func(float64), session *ssh.Session) error {
	f, err := os.Open(src)
	if err != nil {
		return err
	}
	defer f.Close()
	s, err := f.Stat()
	if err != nil {
		return err
	}

	errCh := make(chan error, 1)
	go func() {
		w, err := session.StdinPipe()
		if err != nil {
			errCh <- err
			return
		}
		defer w.Close()
		fmt.Fprintf(w, "C%#o %d %s\n", s.Mode().Perm(), s.Size(), path.Base(src))
		p := &progressWriter{w, 0, s.Size(), progress}
		if _, err := io.Copy(p, f); err != nil {
			errCh <- err
			return
		}
		fmt.Fprint(w, "\x00")
		close(errCh)
	}()

	err = session.Run(fmt.Sprintf("rm -f %s ; scp -t %s", dest, dest))
	select {
	case err := <-errCh:
		return err
	default:
		return err
	}
}

// TODO(benesch): Make progress handling for directories less confusing. The
// SCP protocol makes this challenging, as it does not send the total size of
// all files.
func SCPGet(src, dest string, progress func(float64), session *ssh.Session) error {
	errCh := make(chan error, 1)
	go func() {
		rp, err := session.StdoutPipe()
		if err != nil {
			errCh <- err
			return
		}
		wp, err := session.StdinPipe()
		if err != nil {
			errCh <- err
			return
		}
		defer wp.Close()

		r := bufio.NewReader(rp)
		dirStack := []string{}
		for {
			fmt.Fprint(wp, "\x00")

			line, err := r.ReadBytes('\n')
			if err != nil {
				if len(dirStack) != 0 || err != io.EOF {
					errCh <- err
				}
				close(errCh)
				return
			}

			if line[0] == 'E' {
				dirStack = dirStack[:len(dirStack)-1]
				continue
			}

			var op byte
			var mode uint32
			var size int64
			var name string
			if n, err := fmt.Sscanf(string(line), "%c%o %d %s", &op, &mode, &size, &name); err != nil {
				errCh <- fmt.Errorf("decoding scp directive: %s: %q", err, line)
				return
			} else if n != 4 {
				errCh <- errors.New(string(line))
				return
			}

			fullname := dest
			if len(dirStack) > 0 {
				fullname = filepath.Join(append(dirStack, name)...)
			}

			switch op {
			case 'C':
				f, err := os.Create(fullname)
				if err != nil {
					errCh <- err
					return
				}
				defer f.Close()

				if err := f.Chmod(os.FileMode(mode)); err != nil {
					errCh <- err
					return
				}

				fmt.Fprint(wp, "\x00")

				p := &progressWriter{f, 0, size, progress}
				if _, err := io.Copy(p, io.LimitReader(r, size)); err != nil {
					errCh <- err
					return
				}

				// File data has a trailing null byte.
				_, err = r.ReadByte()
				if err != nil {
					errCh <- err
					return
				}
			case 'D':
				if err := os.Mkdir(fullname, os.FileMode(mode)); err != nil && !os.IsExist(err) {
					errCh <- err
					return
				}
				if len(dirStack) == 0 {
					dirStack = append(dirStack, dest)
				} else {
					dirStack = append(dirStack, name)
				}
			default:
				errCh <- fmt.Errorf("unknown operation '%c'", op)
				return
			}
		}
	}()

	err := session.Run(fmt.Sprintf("scp -qrf %s", src))
	select {
	case err := <-errCh:
		return err
	default:
		return err
	}
}
<commit_msg>Silence warning about encrypted keys<commit_after>package ssh

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/pkg/errors"

	"github.com/cockroachdb/roachprod/config"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/crypto/ssh/knownhosts"
)

var knownHosts ssh.HostKeyCallback
var knownHostsOnce sync.Once
var InsecureIgnoreHostKey bool

func getKnownHosts() ssh.HostKeyCallback {
	knownHostsOnce.Do(func() {
		var err error
		if InsecureIgnoreHostKey {
			knownHosts = ssh.InsecureIgnoreHostKey()
		} else {
			knownHosts, err = knownhosts.New(filepath.Join(os.Getenv("HOME"), ".ssh", "known_hosts"))
			if err != nil {
				log.Fatal(err)
			}
		}
	})
	return knownHosts
}

func getSSHAgentSigners() []ssh.Signer {
	const authSockEnv = "SSH_AUTH_SOCK"
	agentSocket := os.Getenv(authSockEnv)
	if agentSocket == "" {
		return nil
	}
	sock, err := net.Dial("unix", agentSocket)
	if err != nil {
		log.Printf("SSH_AUTH_SOCK set but unable to connect to agent: %s", err)
		return nil
	}
	agent := agent.NewClient(sock)
	signers, err := agent.Signers()
	if err != nil {
		log.Printf("unable to retrieve keys from agent: %s", err)
		return nil
	}
	return signers
}

func getSSHKeySigner(path string, haveAgent bool) ssh.Signer {
	key, err := ioutil.ReadFile(path)
	if err != nil {
		if !os.IsNotExist(err) {
			log.Printf("unable to read SSH key %q: %s", path, err)
		}
		return nil
	}

	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		if strings.Contains(err.Error(), "cannot decode encrypted private key") {
			if !haveAgent {
				log.Printf(
					"skipping encrypted SSH key %q; if necessary, add the key to your SSH agent", path)
			}
		} else {
			log.Printf("unable to parse SSH key %q: %s", path, err)
		}
		return nil
	}
	return signer
}

func getDefaultSSHKeySigners(haveAgent bool) []ssh.Signer {
	var signers []ssh.Signer
	for _, name := range []string{"id_rsa", "google_compute_engine"} {
		s := getSSHKeySigner(filepath.Join(config.OSUser.HomeDir, ".ssh", name), haveAgent)
		if s != nil {
			signers = append(signers, s)
		}
	}
	return signers
}

func newSSHClient(user, host string) (*ssh.Client, net.Conn, error) {
	config := &ssh.ClientConfig{
		User:            user,
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(sshState.signers...)},
		HostKeyCallback: getKnownHosts(),
	}
	config.SetDefaults()

	addr := fmt.Sprintf("%s:22", host)
	conn, err := net.DialTimeout("tcp", addr, 30*time.Second)
	if err != nil {
		return nil, nil, err
	}
	c, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
	if err != nil {
		return nil, nil, err
	}
	return ssh.NewClient(c, chans, reqs), conn, nil
}

type sshClient struct {
	sync.Mutex
	*ssh.Client
}

var sshState = struct {
	signers     []ssh.Signer
	signersInit sync.Once

	clients  map[string]*sshClient
	clientMu sync.Mutex
}{
	clients: map[string]*sshClient{},
}

func NewSSHSession(user, host string) (*ssh.Session, error) {
	sshState.clientMu.Lock()
	target := fmt.Sprintf("%s@%s", user, host)
	client := sshState.clients[target]
	if client == nil {
		client = &sshClient{}
		sshState.clients[target] = client
	}
	sshState.clientMu.Unlock()

	sshState.signersInit.Do(func() {
		sshState.signers = append(sshState.signers, getSSHAgentSigners()...)
		haveAgentSigner := len(sshState.signers) > 0
		sshState.signers = append(sshState.signers, getDefaultSSHKeySigners(haveAgentSigner)...)
	})

	client.Lock()
	defer client.Unlock()
	if client.Client == nil {
		var err error
		client.Client, _, err = newSSHClient(user, host)
		if err != nil {
			return nil, err
		}
	}
	return client.NewSession()
}

func IsSigKill(err error) bool {
	switch t := err.(type) {
	case *ssh.ExitError:
		return t.Signal() == string(ssh.SIGKILL)
	}
	return false
}

type progressWriter struct {
	writer   io.Writer
	done     int64
	total    int64
	progress func(float64)
}

func (p *progressWriter) Write(b []byte) (int, error) {
	n, err := p.writer.Write(b)
	if err == nil {
		p.done += int64(n)
		p.progress(float64(p.done) / float64(p.total))
	}
	return n, err
}

func SCPPut(src, dest string, progress func(float64), session *ssh.Session) error {
	f, err := os.Open(src)
	if err != nil {
		return err
	}
	defer f.Close()
	s, err := f.Stat()
	if err != nil {
		return err
	}

	errCh := make(chan error, 1)
	go func() {
		w, err := session.StdinPipe()
		if err != nil {
			errCh <- err
			return
		}
		defer w.Close()
		fmt.Fprintf(w, "C%#o %d %s\n", s.Mode().Perm(), s.Size(), path.Base(src))
		p := &progressWriter{w, 0, s.Size(), progress}
		if _, err := io.Copy(p, f); err != nil {
			errCh <- err
			return
		}
		fmt.Fprint(w, "\x00")
		close(errCh)
	}()

	err = session.Run(fmt.Sprintf("rm -f %s ; scp -t %s", dest, dest))
	select {
	case err := <-errCh:
		return err
	default:
		return err
	}
}

// TODO(benesch): Make progress handling for directories less confusing. The
// SCP protocol makes this challenging, as it does not send the total size of
// all files.
func SCPGet(src, dest string, progress func(float64), session *ssh.Session) error {
	errCh := make(chan error, 1)
	go func() {
		rp, err := session.StdoutPipe()
		if err != nil {
			errCh <- err
			return
		}
		wp, err := session.StdinPipe()
		if err != nil {
			errCh <- err
			return
		}
		defer wp.Close()

		r := bufio.NewReader(rp)
		dirStack := []string{}
		for {
			fmt.Fprint(wp, "\x00")

			line, err := r.ReadBytes('\n')
			if err != nil {
				if len(dirStack) != 0 || err != io.EOF {
					errCh <- err
				}
				close(errCh)
				return
			}

			if line[0] == 'E' {
				dirStack = dirStack[:len(dirStack)-1]
				continue
			}

			var op byte
			var mode uint32
			var size int64
			var name string
			if n, err := fmt.Sscanf(string(line), "%c%o %d %s", &op, &mode, &size, &name); err != nil {
				errCh <- fmt.Errorf("decoding scp directive: %s: %q", err, line)
				return
			} else if n != 4 {
				errCh <- errors.New(string(line))
				return
			}

			fullname := dest
			if len(dirStack) > 0 {
				fullname = filepath.Join(append(dirStack, name)...)
			}

			switch op {
			case 'C':
				f, err := os.Create(fullname)
				if err != nil {
					errCh <- err
					return
				}
				defer f.Close()

				if err := f.Chmod(os.FileMode(mode)); err != nil {
					errCh <- err
					return
				}

				fmt.Fprint(wp, "\x00")

				p := &progressWriter{f, 0, size, progress}
				if _, err := io.Copy(p, io.LimitReader(r, size)); err != nil {
					errCh <- err
					return
				}

				// File data has a trailing null byte.
				_, err = r.ReadByte()
				if err != nil {
					errCh <- err
					return
				}
			case 'D':
				if err := os.Mkdir(fullname, os.FileMode(mode)); err != nil && !os.IsExist(err) {
					errCh <- err
					return
				}
				if len(dirStack) == 0 {
					dirStack = append(dirStack, dest)
				} else {
					dirStack = append(dirStack, name)
				}
			default:
				errCh <- fmt.Errorf("unknown operation '%c'", op)
				return
			}
		}
	}()

	err := session.Run(fmt.Sprintf("scp -qrf %s", src))
	select {
	case err := <-errCh:
		return err
	default:
		return err
	}
}
<|endoftext|>


<commit_before>package gohub

import (
	"encoding/json"
	"io/ioutil"
	azure "github.com/Azure/azure-sdk-for-go/storage"
	"crypto/md5"
	"encoding/base64"
	"bytes"
	log "github.com/Sirupsen/logrus"
)

const (
	contentMd5 string = `Content-MD5`
)

type StorageSetting struct {
	Name string
	Key string
}

type lease struct {
	Offset string `json:"offset"`
	SeqNo int64 `json:"sequenceNumber"`
	PartitionId string `json:"partitionId"`
	Epoch int64 `json:"epoch"`
	Owner string `json:"owner"`
	Token string `json:"token"`
}

type checkpoint struct {
	offset string
	seqNo int64
}

type storage interface {
	CreateStorage(hub, cg, pid string) error
	GetCheckpoint(hub, cg, pid string) (checkpoint, error)
	SaveCheckpoint(hub, cg, pid string, cp checkpoint) error
}

type azureStorage struct {
	ss StorageSetting
	c azure.Client
}

func newAzureStorage(ss StorageSetting) *azureStorage {
	c, err := azure.NewClient(
		ss.Name,
		ss.Key,
		azure.DefaultBaseURL,
		azure.DefaultAPIVersion,
		false,
	)
	if err != nil {
		panic(err)
	}
	a := &azureStorage{
		c: c,
		ss: ss,
	}
	return a
}

func (s*azureStorage) CreateStorage(hub, cg, pid string) error {
	bs := s.c.GetBlobService()
	_, err := bs.CreateContainerIfNotExists(hub, azure.ContainerAccessTypeBlob)
	if err != nil {
		return err
	}
	bn := getBlobName(cg, pid)
	exists, err := bs.BlobExists(hub, bn)
	if err != nil {
		return err
	}
	if exists {
		log.WithFields(log.Fields{"Hub": hub, "ConsumerGroup": cg, "PartitionID": pid}).Debug("Lease Exists")
		return nil
	}
	if err = bs.CreateBlockBlob(hub, bn); err != nil {
		return err
	}
	if err = s.createNewLease(hub, cg, pid); err != nil {
		return err
	}
	return nil
}

func (s*azureStorage) GetCheckpoint(hub, cg, pid string) (checkpoint, error) {
	l, err := s.getLease(hub, cg, pid)
	if err != nil {
		return checkpoint{}, err
	}
	if l.Offset == "" {
		l.Offset = "-1"
	}
	log.WithFields(log.Fields{
		"Offset": l.Offset,
		"SeqNo": l.SeqNo,
		"PartitionID": l.PartitionId},
	).Debug("Got latest offset")
	return checkpoint{offset: l.Offset, seqNo: l.SeqNo}, nil
}

func (s*azureStorage) SaveCheckpoint(hub, cg, pid string, cp checkpoint) error {
	l, err := s.getLease(hub, cg, pid)
	if err != nil {
		return nil
	}
	l.Offset = cp.offset
	l.SeqNo = cp.seqNo
	if err := s.saveLease(hub, cg, pid, l); err != nil {
		return err
	}
	return nil
}

func (s*azureStorage) getLease(hub, cg, pid string) (lease, error) {
	var l lease
	bs := s.c.GetBlobService()
	r, err := bs.GetBlob(hub, getBlobName(cg, pid))
	if err != nil {
		return l, err
	}
	defer r.Close()
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return l, err
	}
	if err = json.Unmarshal(b, &l); err != nil {
		return l, err
	}
	return l, nil
}

func (s*azureStorage) createNewLease(hub, cg, pid string) error {
	l := lease{}
	l.Offset = "-1"
	l.SeqNo = 0
	l.Epoch = 0
	l.PartitionId = pid
	l.Owner = `gohub`
	l.Token = ``
	if err := s.saveLease(hub, cg, pid, l); err != nil {
		return err
	}
	log.Debug("New Lease Created")
	return nil
}

func (s*azureStorage) saveLease(hub, cg, pid string, l lease) error {
	b, err := json.Marshal(l)
	if err != nil {
		return err
	}

	h := md5.New()
	h.Write(b)
	md5 := base64.StdEncoding.EncodeToString(h.Sum(nil))

	headers := make(map[string]string)
	headers[contentMd5] = md5

	bn := getBlobName(cg, pid)
	bs := s.c.GetBlobService()
	if err := bs.PutBlockWithLength(hub, bn, md5, uint64(len(b)), bytes.NewReader(b), headers); err != nil {
		return err
	}
	if err := bs.PutBlockList(hub, bn, []azure.Block{{ID: md5, Status: azure.BlockStatusLatest}}); err != nil {
		return err
	}
	return nil
}

func getBlobName(cg, pid string) string {
	return cg + "/" + pid
}<commit_msg>Remove unnecessary logs<commit_after>package gohub

import (
	"encoding/json"
	"io/ioutil"
	azure "github.com/Azure/azure-sdk-for-go/storage"
	"crypto/md5"
	"encoding/base64"
	"bytes"
)

const (
	contentMd5 string = `Content-MD5`
)

type StorageSetting struct {
	Name string
	Key string
}

type lease struct {
	Offset string `json:"offset"`
	SeqNo int64 `json:"sequenceNumber"`
	PartitionId string `json:"partitionId"`
	Epoch int64 `json:"epoch"`
	Owner string `json:"owner"`
	Token string `json:"token"`
}

type checkpoint struct {
	offset string
	seqNo int64
}

type storage interface {
	CreateStorage(hub, cg, pid string) error
	GetCheckpoint(hub, cg, pid string) (checkpoint, error)
	SaveCheckpoint(hub, cg, pid string, cp checkpoint) error
}

type azureStorage struct {
	ss StorageSetting
	c azure.Client
}

func newAzureStorage(ss StorageSetting) *azureStorage {
	c, err := azure.NewClient(
		ss.Name,
		ss.Key,
		azure.DefaultBaseURL,
		azure.DefaultAPIVersion,
		false,
	)
	if err != nil {
		panic(err)
	}
	a := &azureStorage{
		c: c,
		ss: ss,
	}
	return a
}

func (s*azureStorage) CreateStorage(hub, cg, pid string) error {
	bs := s.c.GetBlobService()
	_, err := bs.CreateContainerIfNotExists(hub, azure.ContainerAccessTypeBlob)
	if err != nil {
		return err
	}
	bn := getBlobName(cg, pid)
	exists, err := bs.BlobExists(hub, bn)
	if err != nil {
		return err
	}
	if exists {
		return nil
	}
	if err = bs.CreateBlockBlob(hub, bn); err != nil {
		return err
	}
	if err = s.createNewLease(hub, cg, pid); err != nil {
		return err
	}
	return nil
}

func (s*azureStorage) GetCheckpoint(hub, cg, pid string) (checkpoint, error) {
	l, err := s.getLease(hub, cg, pid)
	if err != nil {
		return checkpoint{}, err
	}
	if l.Offset == "" {
		l.Offset = "-1"
	}
	return checkpoint{offset: l.Offset, seqNo: l.SeqNo}, nil
}

func (s*azureStorage) SaveCheckpoint(hub, cg, pid string, cp checkpoint) error {
	l, err := s.getLease(hub, cg, pid)
	if err != nil {
		return nil
	}
	l.Offset = cp.offset
	l.SeqNo = cp.seqNo
	if err := s.saveLease(hub, cg, pid, l); err != nil {
		return err
	}
	return nil
}

func (s*azureStorage) getLease(hub, cg, pid string) (lease, error) {
	var l lease
	bs := s.c.GetBlobService()
	r, err := bs.GetBlob(hub, getBlobName(cg, pid))
	if err != nil {
		return l, err
	}
	defer r.Close()
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return l, err
	}
	if err = json.Unmarshal(b, &l); err != nil {
		return l, err
	}
	return l, nil
}

func (s*azureStorage) createNewLease(hub, cg, pid string) error {
	l := lease{}
	l.Offset = "-1"
	l.SeqNo = 0
	l.Epoch = 0
	l.PartitionId = pid
	l.Owner = `gohub`
	l.Token = ``
	if err := s.saveLease(hub, cg, pid, l); err != nil {
		return err
	}
	return nil
}

func (s*azureStorage) saveLease(hub, cg, pid string, l lease) error {
	b, err := json.Marshal(l)
	if err != nil {
		return err
	}

	h := md5.New()
	h.Write(b)
	md5 := base64.StdEncoding.EncodeToString(h.Sum(nil))

	headers := make(map[string]string)
	headers[contentMd5] = md5

	bn := getBlobName(cg, pid)
	bs := s.c.GetBlobService()
	if err := bs.PutBlockWithLength(hub, bn, md5, uint64(len(b)), bytes.NewReader(b), headers); err != nil {
		return err
	}
	if err := bs.PutBlockList(hub, bn, []azure.Block{{ID: md5, Status: azure.BlockStatusLatest}}); err != nil {
		return err
	}
	return nil
}

func getBlobName(cg, pid string) string {
	return cg + "/" + pid
}<|endoftext|>


<commit_before>package storageconsul

import (
	"context"
	"net"
	"path"
	"strings"
	"time"

	"github.com/caddyserver/certmagic"
	consul "github.com/hashicorp/consul/api"
	"github.com/pteich/errors"
	"go.uber.org/zap"
)

// ConsulStorage allows to store certificates and other TLS resources
// in a shared cluster environment using Consul's key/value-store.
// It uses distributed locks to ensure consistency.
type ConsulStorage struct {
	certmagic.Storage
	ConsulClient *consul.Client
	logger       *zap.SugaredLogger
	locks        map[string]*consul.Lock

	Address     string `json:"address"`
	Token       string `json:"token"`
	Timeout     int    `json:"timeout"`
	Prefix      string `json:"prefix"`
	ValuePrefix string `json:"value_prefix"`
	AESKey      []byte `json:"aes_key"`
	TlsEnabled  bool   `json:"tls_enabled"`
	TlsInsecure bool   `json:"tls_insecure"`
}

// New connects to Consul and returns a ConsulStorage
func New() *ConsulStorage {
	// create ConsulStorage and pre-set values
	s := ConsulStorage{
		locks:       make(map[string]*consul.Lock),
		AESKey:      []byte(DefaultAESKey),
		ValuePrefix: DefaultValuePrefix,
		Prefix:      DefaultPrefix,
		Timeout:     DefaultTimeout,
	}

	return &s
}

func (cs *ConsulStorage) prefixKey(key string) string {
	return path.Join(cs.Prefix, key)
}

// Lock acquires a distributed lock for the given key or blocks until it gets one
func (cs ConsulStorage) Lock(ctx context.Context, key string) error {
	// if we already hold the lock, return early
	if _, exists := cs.locks[key]; exists {
		return nil
	}

	// prepare the lock
	lock, err := cs.ConsulClient.LockKey(cs.prefixKey(key))
	if err != nil {
		return errors.Wrapf(err, "could not create lock for %s", cs.prefixKey(key))
	}

	// acquire the lock and return a channel that is closed upon lost
	lockActive, err := lock.Lock(ctx.Done())
	if err != nil {
		return errors.Wrapf(err, "unable to lock %s", cs.prefixKey(key))
	}

	// auto-unlock and clean list of locks in case of lost
	go func() {
		<-lockActive
		cs.Unlock(key)
	}()

	// save the lock
	cs.locks[key] = lock

	return nil
}

// Unlock releases a specific lock
func (cs ConsulStorage) Unlock(key string) error {
	// check if we own it and unlock
	lock, exists := cs.locks[key]
	if !exists {
		return errors.Errorf("lock %s not found", cs.prefixKey(key))
	}

	err := lock.Unlock()
	if err != nil {
		return errors.Wrapf(err, "unable to unlock %s", cs.prefixKey(key))
	}

	delete(cs.locks, key)
	return nil
}

// Store saves encrypted data value for a key in Consul KV
func (cs ConsulStorage) Store(key string, value []byte) error {
	kv := &consul.KVPair{Key: cs.prefixKey(key)}

	// prepare the stored data
	consulData := &StorageData{
		Value:    value,
		Modified: time.Now(),
	}

	encryptedValue, err := cs.EncryptStorageData(consulData)
	if err != nil {
		return errors.Wrapf(err, "unable to encode data for %s", cs.prefixKey(key))
	}

	kv.Value = encryptedValue

	if _, err = cs.ConsulClient.KV().Put(kv, nil); err != nil {
		return errors.Wrapf(err, "unable to store data for %s", cs.prefixKey(key))
	}

	return nil
}

// Load retrieves the value for a key from Consul KV
func (cs ConsulStorage) Load(key string) ([]byte, error) {
	kv, _, err := cs.ConsulClient.KV().Get(cs.prefixKey(key), &consul.QueryOptions{RequireConsistent: true})
	if err != nil {
		return nil, errors.Wrapf(err, "unable to obtain data for %s", cs.prefixKey(key))
	} else if kv == nil {
		return nil, certmagic.ErrNotExist(errors.Errorf("key %s does not exist", cs.prefixKey(key)))
	}

	contents, err := cs.DecryptStorageData(kv.Value)
	if err != nil {
		return nil, errors.Wrapf(err, "unable to decrypt data for %s", cs.prefixKey(key))
	}

	return contents.Value, nil
}

// Delete a key from Consul KV
func (cs ConsulStorage) Delete(key string) error {
	// first obtain existing keypair
	kv, _, err := cs.ConsulClient.KV().Get(cs.prefixKey(key), &consul.QueryOptions{RequireConsistent: true})
	if err != nil {
		return errors.Wrapf(err, "unable to obtain data for %s", cs.prefixKey(key))
	} else if kv == nil {
		return certmagic.ErrNotExist(err)
	}

	// no do a Check-And-Set operation to verify we really deleted the key
	if success, _, err := cs.ConsulClient.KV().DeleteCAS(kv, nil); err != nil {
		return errors.Wrapf(err, "unable to delete data for %s", cs.prefixKey(key))
	} else if !success {
		return errors.Errorf("failed to lock data delete for %s", cs.prefixKey(key))
	}

	return nil
}

// Exists checks if a key exists
func (cs ConsulStorage) Exists(key string) bool {
	kv, _, err := cs.ConsulClient.KV().Get(cs.prefixKey(key), &consul.QueryOptions{RequireConsistent: true})
	if kv != nil && err == nil {
		return true
	}
	return false
}

// List returns a list with all keys under a given prefix
func (cs ConsulStorage) List(prefix string, recursive bool) ([]string, error) {
	var keysFound []string

	// get a list of all keys at prefix
	keys, _, err := cs.ConsulClient.KV().Keys(cs.prefixKey(prefix), "", &consul.QueryOptions{RequireConsistent: true})
	if err != nil {
		return keysFound, err
	}

	if len(keys) == 0 {
		return keysFound, certmagic.ErrNotExist(errors.Errorf("no keys at %s", prefix))
	}

	// remove default prefix from keys
	for _, key := range keys {
		if strings.HasPrefix(key, cs.prefixKey(prefix)) {
			key = strings.TrimPrefix(key, cs.Prefix+"/")
			keysFound = append(keysFound, key)
		}
	}

	// if recursive wanted, just return all keys
	if recursive {
		return keysFound, nil
	}

	// for non-recursive split path and look for unique keys just under given prefix
	keysMap := make(map[string]bool)
	for _, key := range keysFound {
		dir := strings.Split(strings.TrimPrefix(key, prefix+"/"), "/")
		keysMap[dir[0]] = true
	}

	keysFound = make([]string, 0)
	for key := range keysMap {
		keysFound = append(keysFound, path.Join(prefix, key))
	}

	return keysFound, nil
}

// Stat returns statistic data of a key
func (cs ConsulStorage) Stat(key string) (certmagic.KeyInfo, error) {
	kv, _, err := cs.ConsulClient.KV().Get(cs.prefixKey(key), &consul.QueryOptions{RequireConsistent: true})
	if err != nil {
		return certmagic.KeyInfo{}, errors.Errorf("unable to obtain data for %s", cs.prefixKey(key))
	} else if kv == nil {
		return certmagic.KeyInfo{}, certmagic.ErrNotExist(errors.Errorf("key %s does not exist", cs.prefixKey(key)))
	}

	contents, err := cs.DecryptStorageData(kv.Value)
	if err != nil {
		return certmagic.KeyInfo{}, errors.Errorf("unable to decrypt data for %s", cs.prefixKey(key))
	}

	return certmagic.KeyInfo{
		Key:        key,
		Modified:   contents.Modified,
		Size:       int64(len(contents.Value)),
		IsTerminal: false,
	}, nil
}

func (cs *ConsulStorage) createConsulClient() error {
	// get the default config
	consulCfg := consul.DefaultConfig()
	if cs.Address != "" {
		consulCfg.Address = cs.Address
	}
	if cs.Token != "" {
		consulCfg.Token = cs.Token
	}
	if cs.TlsEnabled {
		consulCfg.Scheme = "https"
	}
	consulCfg.TLSConfig.InsecureSkipVerify = cs.TlsInsecure

	// set a dial context to prevent default keepalive
	consulCfg.Transport.DialContext = (&net.Dialer{
		Timeout:   time.Duration(cs.Timeout) * time.Second,
		KeepAlive: time.Duration(cs.Timeout) * time.Second,
	}).DialContext

	// create the Consul API client
	consulClient, err := consul.NewClient(consulCfg)
	if err != nil {
		return errors.Wrap(err, "unable to create Consul client")
	}
	if _, err := consulClient.Agent().NodeName(); err != nil {
		return errors.Wrap(err, "unable to ping Consul")
	}

	cs.ConsulClient = consulClient
	return nil
}
<commit_msg>Add mutex to protect local locks map<commit_after>package storageconsul

import (
	"context"
	"net"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/caddyserver/certmagic"
	consul "github.com/hashicorp/consul/api"
	"github.com/pteich/errors"
	"go.uber.org/zap"
)

// ConsulStorage allows to store certificates and other TLS resources
// in a shared cluster environment using Consul's key/value-store.
// It uses distributed locks to ensure consistency.
type ConsulStorage struct {
	certmagic.Storage
	ConsulClient *consul.Client
	logger       *zap.SugaredLogger
	muLocks      sync.Mutex
	locks        map[string]*consul.Lock

	Address     string `json:"address"`
	Token       string `json:"token"`
	Timeout     int    `json:"timeout"`
	Prefix      string `json:"prefix"`
	ValuePrefix string `json:"value_prefix"`
	AESKey      []byte `json:"aes_key"`
	TlsEnabled  bool   `json:"tls_enabled"`
	TlsInsecure bool   `json:"tls_insecure"`
}

// New connects to Consul and returns a ConsulStorage
func New() *ConsulStorage {
	// create ConsulStorage and pre-set values
	s := ConsulStorage{
		locks:       make(map[string]*consul.Lock),
		AESKey:      []byte(DefaultAESKey),
		ValuePrefix: DefaultValuePrefix,
		Prefix:      DefaultPrefix,
		Timeout:     DefaultTimeout,
	}

	return &s
}

func (cs *ConsulStorage) prefixKey(key string) string {
	return path.Join(cs.Prefix, key)
}

// Lock acquires a distributed lock for the given key or blocks until it gets one
func (cs ConsulStorage) Lock(ctx context.Context, key string) error {
	cs.muLocks.Lock()
	defer cs.muLocks.Unlock()

	// if we already hold the lock, return early
	if _, exists := cs.locks[key]; exists {
		return nil
	}

	// prepare the lock
	lock, err := cs.ConsulClient.LockKey(cs.prefixKey(key))
	if err != nil {
		return errors.Wrapf(err, "could not create lock for %s", cs.prefixKey(key))
	}

	// acquire the lock and return a channel that is closed upon lost
	lockActive, err := lock.Lock(ctx.Done())
	if err != nil {
		return errors.Wrapf(err, "unable to lock %s", cs.prefixKey(key))
	}

	// auto-unlock and clean list of locks in case of lost
	go func() {
		<-lockActive
		cs.Unlock(key)
	}()

	// save the lock
	cs.muLocks.Lock()
	cs.locks[key] = lock
	cs.muLocks.Unlock()

	return nil
}

// Unlock releases a specific lock
func (cs ConsulStorage) Unlock(key string) error {
	cs.muLocks.Lock()
	defer cs.muLocks.Unlock()

	// check if we own it and unlock
	lock, exists := cs.locks[key]
	if !exists {
		return errors.Errorf("lock %s not found", cs.prefixKey(key))
	}

	err := lock.Unlock()
	if err != nil {
		return errors.Wrapf(err, "unable to unlock %s", cs.prefixKey(key))
	}

	delete(cs.locks, key)
	return nil
}

// Store saves encrypted data value for a key in Consul KV
func (cs ConsulStorage) Store(key string, value []byte) error {
	kv := &consul.KVPair{Key: cs.prefixKey(key)}

	// prepare the stored data
	consulData := &StorageData{
		Value:    value,
		Modified: time.Now(),
	}

	encryptedValue, err := cs.EncryptStorageData(consulData)
	if err != nil {
		return errors.Wrapf(err, "unable to encode data for %s", cs.prefixKey(key))
	}

	kv.Value = encryptedValue

	if _, err = cs.ConsulClient.KV().Put(kv, nil); err != nil {
		return errors.Wrapf(err, "unable to store data for %s", cs.prefixKey(key))
	}

	return nil
}

// Load retrieves the value for a key from Consul KV
func (cs ConsulStorage) Load(key string) ([]byte, error) {
	kv, _, err := cs.ConsulClient.KV().Get(cs.prefixKey(key), &consul.QueryOptions{RequireConsistent: true})
	if err != nil {
		return nil, errors.Wrapf(err, "unable to obtain data for %s", cs.prefixKey(key))
	} else if kv == nil {
		return nil, certmagic.ErrNotExist(errors.Errorf("key %s does not exist", cs.prefixKey(key)))
	}

	contents, err := cs.DecryptStorageData(kv.Value)
	if err != nil {
		return nil, errors.Wrapf(err, "unable to decrypt data for %s", cs.prefixKey(key))
	}

	return contents.Value, nil
}

// Delete a key from Consul KV
func (cs ConsulStorage) Delete(key string) error {
	// first obtain existing keypair
	kv, _, err := cs.ConsulClient.KV().Get(cs.prefixKey(key), &consul.QueryOptions{RequireConsistent: true})
	if err != nil {
		return errors.Wrapf(err, "unable to obtain data for %s", cs.prefixKey(key))
	} else if kv == nil {
		return certmagic.ErrNotExist(err)
	}

	// no do a Check-And-Set operation to verify we really deleted the key
	if success, _, err := cs.ConsulClient.KV().DeleteCAS(kv, nil); err != nil {
		return errors.Wrapf(err, "unable to delete data for %s", cs.prefixKey(key))
	} else if !success {
		return errors.Errorf("failed to lock data delete for %s", cs.prefixKey(key))
	}

	return nil
}

// Exists checks if a key exists
func (cs ConsulStorage) Exists(key string) bool {
	kv, _, err := cs.ConsulClient.KV().Get(cs.prefixKey(key), &consul.QueryOptions{RequireConsistent: true})
	if kv != nil && err == nil {
		return true
	}
	return false
}

// List returns a list with all keys under a given prefix
func (cs ConsulStorage) List(prefix string, recursive bool) ([]string, error) {
	var keysFound []string

	// get a list of all keys at prefix
	keys, _, err := cs.ConsulClient.KV().Keys(cs.prefixKey(prefix), "", &consul.QueryOptions{RequireConsistent: true})
	if err != nil {
		return keysFound, err
	}

	if len(keys) == 0 {
		return keysFound, certmagic.ErrNotExist(errors.Errorf("no keys at %s", prefix))
	}

	// remove default prefix from keys
	for _, key := range keys {
		if strings.HasPrefix(key, cs.prefixKey(prefix)) {
			key = strings.TrimPrefix(key, cs.Prefix+"/")
			keysFound = append(keysFound, key)
		}
	}

	// if recursive wanted, just return all keys
	if recursive {
		return keysFound, nil
	}

	// for non-recursive split path and look for unique keys just under given prefix
	keysMap := make(map[string]bool)
	for _, key := range keysFound {
		dir := strings.Split(strings.TrimPrefix(key, prefix+"/"), "/")
		keysMap[dir[0]] = true
	}

	keysFound = make([]string, 0)
	for key := range keysMap {
		keysFound = append(keysFound, path.Join(prefix, key))
	}

	return keysFound, nil
}

// Stat returns statistic data of a key
func (cs ConsulStorage) Stat(key string) (certmagic.KeyInfo, error) {
	kv, _, err := cs.ConsulClient.KV().Get(cs.prefixKey(key), &consul.QueryOptions{RequireConsistent: true})
	if err != nil {
		return certmagic.KeyInfo{}, errors.Errorf("unable to obtain data for %s", cs.prefixKey(key))
	} else if kv == nil {
		return certmagic.KeyInfo{}, certmagic.ErrNotExist(errors.Errorf("key %s does not exist", cs.prefixKey(key)))
	}

	contents, err := cs.DecryptStorageData(kv.Value)
	if err != nil {
		return certmagic.KeyInfo{}, errors.Errorf("unable to decrypt data for %s", cs.prefixKey(key))
	}

	return certmagic.KeyInfo{
		Key:        key,
		Modified:   contents.Modified,
		Size:       int64(len(contents.Value)),
		IsTerminal: false,
	}, nil
}

func (cs *ConsulStorage) createConsulClient() error {
	// get the default config
	consulCfg := consul.DefaultConfig()
	if cs.Address != "" {
		consulCfg.Address = cs.Address
	}
	if cs.Token != "" {
		consulCfg.Token = cs.Token
	}
	if cs.TlsEnabled {
		consulCfg.Scheme = "https"
	}
	consulCfg.TLSConfig.InsecureSkipVerify = cs.TlsInsecure

	// set a dial context to prevent default keepalive
	consulCfg.Transport.DialContext = (&net.Dialer{
		Timeout:   time.Duration(cs.Timeout) * time.Second,
		KeepAlive: time.Duration(cs.Timeout) * time.Second,
	}).DialContext

	// create the Consul API client
	consulClient, err := consul.NewClient(consulCfg)
	if err != nil {
		return errors.Wrap(err, "unable to create Consul client")
	}
	if _, err := consulClient.Agent().NodeName(); err != nil {
		return errors.Wrap(err, "unable to ping Consul")
	}

	cs.ConsulClient = consulClient
	return nil
}
<|endoftext|>


<commit_before>package dockerface

import (
	"fmt"
	"net/http"
	"time"

	"github.com/ant0ine/go-json-rest/rest"
	"github.com/jmcvetta/napping"
)

type ContainerEndpoint struct{}

func (ce *ContainerEndpoint) Filter(w rest.ResponseWriter, req *rest.Request) {
	params := napping.Params{"all": req.URL.Query().Get("all")}

	var result []*Container
	resp, err := napping.Get(buildUrl("containers/json"), &params, &result, nil)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if resp.Status() != http.StatusOK {
		rest.Error(w, fmt.Sprintf("Invalid status code: %v", resp.Status()), http.StatusBadRequest)
		return
	}
	w.WriteJson(result)
}

func (ce *ContainerEndpoint) Start(w rest.ResponseWriter, req *rest.Request) {
	cid := req.PathParam("id")
	resp, err := napping.Post(buildUrl("containers/"+cid+"/start"), nil, nil, nil)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// if start was successful, retrieve container and return its current state
	if resp.Status() == http.StatusNoContent {
		time.Sleep(1000 * time.Millisecond) // give it time

		var result []*Container
		resp, err := napping.Get(buildUrl("containers/json"), nil, &result, nil)
		if err != nil {
			rest.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if resp.Status() == http.StatusOK {
			if c := filterContainer(result, cid); c != nil {
				w.WriteJson(c)
			}
		}
	} else {
		w.WriteHeader(resp.Status())
	}
}

func (ce *ContainerEndpoint) Stop(w rest.ResponseWriter, req *rest.Request) {
	cid := req.PathParam("id")
	resp, err := napping.Post(buildUrl("containers/"+cid+"/stop"), nil, nil, nil)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// if stop was successful, retrieve container and return its current state
	if resp.Status() == http.StatusNoContent {
		time.Sleep(1000 * time.Millisecond) // give it time

		params := napping.Params{"all": "true"}
		var result []*Container
		resp, err := napping.Get(buildUrl("containers/json"), &params, &result, nil)
		if err != nil {
			rest.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if resp.Status() == http.StatusOK {
			if c := filterContainer(result, cid); c != nil {
				w.WriteJson(c)
			}
		}
	} else {
		w.WriteHeader(resp.Status())
	}
}

func filterContainer(containers []*Container, cid string) *Container {
	for _, c := range containers {
		if c.Id == cid {
			return c
		}
	}
	return nil
}
<commit_msg>optimize returning of state when starting/stopping<commit_after>package dockerface

import (
	"fmt"
	"net/http"

	"github.com/ant0ine/go-json-rest/rest"
	"github.com/jmcvetta/napping"
)

type stateContainer struct {
	State struct {
		Paused     bool   `json:"paused"`
		Restarting bool   `json:"restarting"`
		Running    bool   `json:"running"`
		FinishedAt string `json:"finishedat"`
		StartedAt  string `json:"startedat"`
		ExitCode   int    `json:"exitcode"`
	} `json:"state"`
}

type ContainerEndpoint struct{}

func (ce *ContainerEndpoint) Filter(w rest.ResponseWriter, req *rest.Request) {
	params := napping.Params{"all": req.URL.Query().Get("all")}

	var result []*Container
	resp, err := napping.Get(buildUrl("containers/json"), &params, &result, nil)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if resp.Status() != http.StatusOK {
		rest.Error(w, fmt.Sprintf("Invalid status code: %v", resp.Status()), http.StatusBadRequest)
		return
	}
	w.WriteJson(result)
}

func (ce *ContainerEndpoint) Start(w rest.ResponseWriter, req *rest.Request) {
	cid := req.PathParam("id")
	resp, err := napping.Post(buildUrl("containers/"+cid+"/start"), nil, nil, nil)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// if start was successful, retrieve container and return its current state
	if resp.Status() == http.StatusNoContent {
		container, err := getContainerState(cid)
		if err != nil {
			rest.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.WriteJson(container)
		}
	} else {
		w.WriteHeader(resp.Status())
	}
}

func (ce *ContainerEndpoint) Stop(w rest.ResponseWriter, req *rest.Request) {
	cid := req.PathParam("id")
	resp, err := napping.Post(buildUrl("containers/"+cid+"/stop"), nil, nil, nil)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// if stop was successful, retrieve container and return its current state
	if resp.Status() == http.StatusNoContent {
		container, err := getContainerState(cid)
		if err != nil {
			rest.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.WriteJson(container)
		}
	} else {
		w.WriteHeader(resp.Status())
	}
}

func getContainerState(cid string) (*stateContainer, error) {
	var container *stateContainer
	_, err := napping.Get(buildUrl("containers/"+cid+"/json"), nil, &container, nil)
	if err != nil {
		return nil, fmt.Errorf("Error getting container state: %v", err)
	}
	return container, nil
}
<|endoftext|>


<commit_before>package imp

import "github.com/qlova/ilang/src"
import "os"
import "text/scanner"

func init() {
	ilang.RegisterToken([]string{"import"}, ScanImport)
}

//Messy.
func ScanImport(ic *ilang.Compiler) {
	pkg := ic.Scan(ilang.Name)
	ic.Scan('\n')

	var filename = ""

	retry:
	file, err := os.Open(pkg+".i")
	if err != nil {
		if file, err = os.Open(pkg+"/"+pkg+".i"); err != nil {
			dir, _ := os.Getwd()
			if ic.FileDepth > 0 {
				ic.FileDepth--
				os.Chdir(ic.Dirs[len(ic.Dirs)-1])
				ic.Dirs = ic.Dirs[:len(ic.Dirs)-1]
				goto retry
			}

			ic.RaiseError("Cannot import "+pkg+", does not exist!", dir)
		} else {
			 filename = pkg+"/"+pkg+".i"

			 dir, _ := os.Getwd()

			ic.Dirs = append(ic.Dirs, dir)

			os.Chdir("./"+pkg)
			ic.FileDepth++
		}
	} else {
		filename = pkg+".i"
	}
	ic.Scanners = append(ic.Scanners, ic.Scanner)

	ic.Scanner = &scanner.Scanner{}
	ic.Scanner.Init(file)
	ic.Scanner.Position.Filename = filename
	ic.Scanner.Whitespace= 1<<'\t' | 1<<'\r' | 1<<' '
}
<commit_msg>Now, libraries can be loaded from ~/.ilang/imports/...<commit_after>package imp

import "github.com/qlova/ilang/src"
import "os"
import "text/scanner"
import "os/user"

func init() {
	ilang.RegisterToken([]string{"import"}, ScanImport)
}

//Messy.
func ScanImport(ic *ilang.Compiler) {
	pkg := ic.Scan(ilang.Name)
	ic.Scan('\n')

	var filename = ""

	retry:
	file, err := os.Open(pkg+".i")
	if err != nil {
		if file, err = os.Open(pkg+"/"+pkg+".i"); err != nil {

			//Search through parent folders?
			dir, _ := os.Getwd()
			if ic.FileDepth > 0 {
				ic.FileDepth--
				os.Chdir(ic.Dirs[len(ic.Dirs)-1])
				ic.Dirs = ic.Dirs[:len(ic.Dirs)-1]
				goto retry
			}

			//Search in ~/.ilang.
			if ic.FileDepth == 0 {
				usr, err := user.Current()
				if err == nil {
					os.Chdir(usr.HomeDir+"/.ilang/imports/")
					ic.Dirs = append(ic.Dirs, usr.HomeDir+"/.ilang/imports/")
					goto retry
				}
			}

			ic.RaiseError("Cannot import "+pkg+", does not exist!", dir)
		} else {
			 filename = pkg+"/"+pkg+".i"

			 dir, _ := os.Getwd()

			ic.Dirs = append(ic.Dirs, dir)

			os.Chdir("./"+pkg)
			ic.FileDepth++
		}
	} else {
		filename = pkg+".i"
	}
	ic.Scanners = append(ic.Scanners, ic.Scanner)

	ic.Scanner = &scanner.Scanner{}
	ic.Scanner.Init(file)
	ic.Scanner.Position.Filename = filename
	ic.Scanner.Whitespace= 1<<'\t' | 1<<'\r' | 1<<' '
}
<|endoftext|>


<commit_before>package file

import (
	"appengine"
	"appengine/datastore"
	"appengine/file"

	"encoding/json"
	"io"
	"log"
	"mime/multipart"
	"net/http"
	"time"

	"code.google.com/p/go-uuid/uuid"
)

type BlobContent struct {
	Id          string    `json: "id" datastore:"_"`
	Filename    string    `json: "filename"`
	AbsFilename string    `json: "absFilename"`
	ContentType string    `json: "contentType"`
	Size        int64     `json: "size"`
	CreatedAt   time.Time `json: "createdAt"`
}

func init() {
	http.HandleFunc("/file", handler)
}

func handler(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)

	switch r.Method {
	default:
		http.Error(w, "unsupported method.", http.StatusMethodNotAllowed)
	case "POST":
		uploadFile(c, w, r)
	case "GET":
		downloadFile(c, w, r)
	}
}

func uploadFile(c appengine.Context, w http.ResponseWriter, r *http.Request) {
	err := r.ParseMultipartForm(10 * 1024 * 1024) // 10MB
	if err != nil {
		if err.Error() == "permission denied" {
			http.Error(w, "Upload the file is too large.\n", http.StatusBadRequest)
			return
		} else {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
	f, fh, err := r.FormFile("filename")
	if err != nil {
		c.Errorf("%s", err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer f.Close()

	log.Printf("%s", fh.Filename)
	log.Printf("%v", fh.Header)
	log.Printf("Content-Type : %s", fh.Header.Get("Content-Type"))

	absFilename, size, err := directStore(c, f, fh)
	if err != nil {
		c.Errorf("%s", err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	log.Printf("absFilename : %s", absFilename)
	log.Printf("size : %d", size)

	b := &BlobContent{
		uuid.New(),
		fh.Filename,
		absFilename,
		fh.Header.Get("Content-Type"),
		size,
		time.Now(),
	}
	_, err = b.Save(c)
	if err != nil {
		c.Errorf("%s", err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
	err = json.NewEncoder(w).Encode(b)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func downloadFile(c appengine.Context, w http.ResponseWriter, r *http.Request) {
	bn, err := file.DefaultBucketName(c)
	if err != nil {
		c.Errorf("%s", err.Error())
	}

	// JSTの日ごとにPathを区切っておく
	filename := "/gs/" + bn + "/" + r.FormValue("name")
	log.Printf("filename : " + filename)
	fr, err := file.Open(c, filename)
	if err != nil {
		c.Errorf("%s", err.Error())
	}
	defer fr.Close()

	w.Header().Set("Cache-Control:public", "max-age=120")
	io.Copy(w, fr)
}

func directStore(c appengine.Context, f multipart.File, fh *multipart.FileHeader) (absFilename string, size int64, err error) {
	bn, err := file.DefaultBucketName(c)
	if err != nil {
		return "", 0, err
	}

	opts := &file.CreateOptions{
		MIMEType:   fh.Header.Get("Content-Type"),
		BucketName: bn,
	}

	wc, absFilename, err := file.Create(c, getNowDateJst(time.Now())+"/"+uuid.New(), opts)
	if err != nil {
		return "", 0, err
	}
	defer wc.Close()

	size, err = io.Copy(wc, f)
	if err != nil {
		return "", 0, err
	}

	return absFilename, size, nil
}

func getNowDateJst(t time.Time) string {
	j := t.In(time.FixedZone("Asia/Tokyo", 9*60*60))
	return j.Format("20060102")
}

func CreateBlobContentKey(c appengine.Context, id string) *datastore.Key {
	log.Printf("key name = %s : ", id)
	return datastore.NewKey(c, "BlobContent", id, 0, nil)
}

func (b *BlobContent) Key(c appengine.Context) *datastore.Key {
	return CreateBlobContentKey(c, b.Id)
}

func (b *BlobContent) Save(c appengine.Context) (*BlobContent, error) {
	b.CreatedAt = time.Now()
	k, err := datastore.Put(c, b.Key(c), b)
	if err != nil {
		return nil, err
	}

	b.Id = k.StringID()
	return b, nil
}
<commit_msg>feat(gcs) : add Response Content-Type<commit_after>package file

import (
	"appengine"
	"appengine/datastore"
	"appengine/file"

	"encoding/json"
	"io"
	"log"
	"mime/multipart"
	"net/http"
	"time"

	"code.google.com/p/go-uuid/uuid"
)

type BlobContent struct {
	Id          string    `json: "id" datastore:"_"`
	Filename    string    `json: "filename"`
	AbsFilename string    `json: "absFilename"`
	ContentType string    `json: "contentType"`
	Size        int64     `json: "size"`
	CreatedAt   time.Time `json: "createdAt"`
}

func init() {
	http.HandleFunc("/file", handler)
}

func handler(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)

	switch r.Method {
	default:
		http.Error(w, "unsupported method.", http.StatusMethodNotAllowed)
	case "POST":
		uploadFile(c, w, r)
	case "GET":
		downloadFile(c, w, r)
	}
}

func uploadFile(c appengine.Context, w http.ResponseWriter, r *http.Request) {
	err := r.ParseMultipartForm(10 * 1024 * 1024) // 10MB
	if err != nil {
		if err.Error() == "permission denied" {
			http.Error(w, "Upload the file is too large.\n", http.StatusBadRequest)
			return
		} else {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
	f, fh, err := r.FormFile("filename")
	if err != nil {
		c.Errorf("%s", err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer f.Close()

	log.Printf("%s", fh.Filename)
	log.Printf("%v", fh.Header)
	log.Printf("Content-Type : %s", fh.Header.Get("Content-Type"))

	absFilename, size, err := directStore(c, f, fh)
	if err != nil {
		c.Errorf("%s", err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	log.Printf("absFilename : %s", absFilename)
	log.Printf("size : %d", size)

	b := &BlobContent{
		uuid.New(),
		fh.Filename,
		absFilename,
		fh.Header.Get("Content-Type"),
		size,
		time.Now(),
	}
	_, err = b.Save(c)
	if err != nil {
		c.Errorf("%s", err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
	err = json.NewEncoder(w).Encode(b)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

func downloadFile(c appengine.Context, w http.ResponseWriter, r *http.Request) {
	id := r.FormValue("id")
	k := CreateBlobContentKey(c, id)

	var b BlobContent
	err := datastore.Get(c, k, &b)
	if err != nil {
		c.Errorf("%s", err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	fr, err := file.Open(c, b.AbsFilename)
	if err != nil {
		c.Errorf("%s", err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer fr.Close()

	w.Header().Set("Cache-Control:public", "max-age=120")
	w.Header().Set("Content-Type", b.ContentType)
	io.Copy(w, fr)
}

func directStore(c appengine.Context, f multipart.File, fh *multipart.FileHeader) (absFilename string, size int64, err error) {
	bn, err := file.DefaultBucketName(c)
	if err != nil {
		return "", 0, err
	}

	opts := &file.CreateOptions{
		MIMEType:   fh.Header.Get("Content-Type"),
		BucketName: bn,
	}

	// JSTで、日ごとにPathを区切っておく
	wc, absFilename, err := file.Create(c, getNowDateJst(time.Now())+"/"+uuid.New(), opts)
	if err != nil {
		return "", 0, err
	}
	defer wc.Close()

	size, err = io.Copy(wc, f)
	if err != nil {
		return "", 0, err
	}

	return absFilename, size, nil
}

func getNowDateJst(t time.Time) string {
	j := t.In(time.FixedZone("Asia/Tokyo", 9*60*60))
	return j.Format("20060102")
}

func CreateBlobContentKey(c appengine.Context, id string) *datastore.Key {
	log.Printf("key name = %s : ", id)
	return datastore.NewKey(c, "BlobContent", id, 0, nil)
}

func (b *BlobContent) Key(c appengine.Context) *datastore.Key {
	return CreateBlobContentKey(c, b.Id)
}

func (b *BlobContent) Save(c appengine.Context) (*BlobContent, error) {
	b.CreatedAt = time.Now()
	k, err := datastore.Put(c, b.Key(c), b)
	if err != nil {
		return nil, err
	}

	b.Id = k.StringID()
	return b, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Davis Webb
// Copyright 2015 Guntas Grewal
// Copyright 2015 Luke Shumaker

package store

import (
	"database/sql"
	"golang.org/x/crypto/bcrypt"
	he "httpentity"
	"strings"
)

var _ he.Entity = &User{}
var _ he.NetEntity = &User{}
var dirUsers he.Entity = newDirUsers()

// Model /////////////////////////////////////////////////////////////

type User struct {
	Id       string
	FullName string
	Email    string
	pwHash   []byte
}

func getUserById(con DB, id int) *User {
	var user User
	err := con.QueryRow("SELECT * FROM users WHERE id=?", id).Scan(&user)
	switch {
	case err == sql.ErrNoRows:
		// user does not exist
		return nil
	case err != nil:
		// error talking to the DB
		panic(err)
	default:
		// all ok
		return &user
	}
}

func GetUserByName(con DB, name string) *User {
	var user User
	err := con.QueryRow("SELECT * FROM users WHERE name=?", name).Scan(&user)
	switch {
	case err == sql.ErrNoRows:
		// user does not exist
		return nil
	case err != nil:
		// error talking to the DB
		panic(err)
	default:
		// all ok
		return &user
	}
}

func GetUserByEmail(con DB, address string) *User {
	var user User
	err := con.QueryRow(""+
		"SELECT users.* "+
		"FROM users JOIN user_address ON users.id=user_address.user_id "+
		"WHERE user_address.address=? AND user_address=?",
		address,
		"email").Scan(&user)
	switch {
	case err == sql.ErrNoRows:
		// user does not exist
		return nil
	case err != nil:
		// error talking to the DB
		panic(err)
	default:
		// all ok
		return &user
	}
}

func (u *User) SetPassword(password string) error {
	hash, err := bcrypt.GenerateFromPassword([]byte(password), -1)
	u.pwHash = hash
	return err
}

func (u *User) CheckPassword(password string) bool {
	err := bcrypt.CompareHashAndPassword(u.pwHash, []byte(password))
	return err != nil
}

func NewUser(con DB, name string, password string, email string) *User {
	u := &User{
		Id:       name,
		FullName: "",
		Email:    email,
	}
	u.SetPassword(password)
	_, err := con.Exec("INSERT INTO users VALUES (?,?,?,?)", u.Id, u.FullName, u.pwHash, u.Email)
	if err != nil {
		panic(err)
	}
	return u
}

func (u *User) Save() {
	dbMap.Update(u)
}

func (o *User) Subentity(name string, req he.Request) he.Entity {
	return nil
}

func (o *User) Methods() map[string]he.Handler {
	return map[string]he.Handler{
		"GET": func(req he.Request) he.Response {
			panic("TODO: API: (*User).Methods()[\"GET\"]")
		},
		"PUT": func(req he.Request) he.Response {
			panic("TODO: API: (*User).Methods()[\"PUT\"]")
		},
		"PATCH": func(req he.Request) he.Response {
			panic("TODO: API: (*User).Methods()[\"PATCH\"]")
		},
		"DELETE": func(req he.Request) he.Response {
			panic("TODO: API: (*User).Methods()[\"DELETE\"]")
		},
	}
}

// View //////////////////////////////////////////////////////////////

func (o *User) Encoders() map[string]he.Encoder {
	return defaultEncoders(o)
}

// Directory ("Controller") //////////////////////////////////////////

type t_dirUsers struct {
	methods map[string]he.Handler
}

func newDirUsers() t_dirUsers {
	r := t_dirUsers{}
	r.methods = map[string]he.Handler{
		"POST": func(req he.Request) he.Response {
			db := req.Things["db"].(DB)
			badbody := req.StatusBadRequest("submitted body not what expected")
			hash, ok := req.Entity.(map[string]interface{}); if !ok { return badbody }
			username, ok := hash["username"].(string)      ; if !ok { return badbody }
			email   , ok := hash["email"].(string)         ; if !ok { return badbody }
			password, ok := hash["password"].(string)      ; if !ok { return badbody }

			if password2, ok := hash["password_verification"].(string); ok {
				if password != password2 {
					// Passwords don't match
					return req.StatusConflict(he.NetString("password and password_verification don't match"))
				}
			}

			username = strings.ToLower(username)

			user := NewUser(db, username, password, email)
			if user == nil {
				return req.StatusConflict(he.NetString("either that username or password is already taken"))
			} else {
				return req.StatusCreated(r, username)
			}
		},
	}
	return r
}

func (d t_dirUsers) Methods() map[string]he.Handler {
	return d.methods
}

func (d t_dirUsers) Subentity(name string, req he.Request) he.Entity {
	db := req.Things["db"].(DB)
	return GetUserByName(db, name)
}
<commit_msg>added user api's<commit_after>// Copyright 2015 Davis Webb
// Copyright 2015 Guntas Grewal
// Copyright 2015 Luke Shumaker

package store

import (
	"database/sql"
	"golang.org/x/crypto/bcrypt"
	he "httpentity"
	"strings"
)

var _ he.Entity = &User{}
var _ he.NetEntity = &User{}
var dirUsers he.Entity = newDirUsers()

// Model /////////////////////////////////////////////////////////////

type User struct {
	Id       string
	FullName string
	Email    string
	pwHash   []byte
}

func getUserById(con DB, id int) *User {
	var user User
	err := con.QueryRow("SELECT * FROM users WHERE id=?", id).Scan(&user)
	switch {
	case err == sql.ErrNoRows:
		// user does not exist
		return nil
	case err != nil:
		// error talking to the DB
		panic(err)
	default:
		// all ok
		return &user
	}
}

func GetUserByName(con DB, name string) *User {
	var user User
	err := con.QueryRow("SELECT * FROM users WHERE name=?", name).Scan(&user)
	switch {
	case err == sql.ErrNoRows:
		// user does not exist
		return nil
	case err != nil:
		// error talking to the DB
		panic(err)
	default:
		// all ok
		return &user
	}
}

func GetUserByEmail(con DB, address string) *User {
	var user User
	err := con.QueryRow(""+
		"SELECT users.* "+
		"FROM users JOIN user_address ON users.id=user_address.user_id "+
		"WHERE user_address.address=? AND user_address=?",
		address,
		"email").Scan(&user)
	switch {
	case err == sql.ErrNoRows:
		// user does not exist
		return nil
	case err != nil:
		// error talking to the DB
		panic(err)
	default:
		// all ok
		return &user
	}
}

func (u *User) SetPassword(password string) error {
	hash, err := bcrypt.GenerateFromPassword([]byte(password), -1)
	u.pwHash = hash
	return err
}

func (u *User) CheckPassword(password string) bool {
	err := bcrypt.CompareHashAndPassword(u.pwHash, []byte(password))
	return err != nil
}

func NewUser(con DB, name string, password string, email string) *User {
	u := &User{
		Id:       name,
		FullName: "",
		Email:    email,
	}
	u.SetPassword(password)
	_, err := con.Exec("INSERT INTO users VALUES (?,?,?,?)", u.Id, u.FullName, u.pwHash, u.Email)
	if err != nil {
		panic(err)
	}
	return u
}

func (u *User) Save() {
	dbMap.Update(u)
}

func (o *User) Subentity(name string, req he.Request) he.Entity {
	return nil
}

func (o *User) Methods() map[string]he.Handler {
	return map[string]he.Handler{
		"GET": func(req he.Request) he.Response {
			sess, ok := req.Entity.Things["session"]; if !ok { return badbody }
			if sess.UserId !=  o.Id {
				return req.StatusUnauthorized(o)
			}
			else {
				req.Entity.Things["user"], ok = getUserById(req.Entity.Things["db"], sess.UserId) ; if !ok { return badbody }
				return req.StatusOK(o)
			}
		},
		"PUT": func(req he.Request) he.Response {
			panic("TODO: API: (*User).Methods()[\"PUT\"]")
		},
		"PATCH": func(req he.Request) he.Response {
                        badbody := req.StatusBadRequest("submitted body not what expected")
                        hash, ok := req.Entity.(map[string]interface{}); if !ok { return badbody }
                        email   , ok := hash["email"].(string)         ; if !ok { return badbody }
			password, ok := hash["password"].(string)      ; if !ok { return badbody }

			if password2, ok := hash["password_verification"].(string); ok {
                                if password != password2 {
                                        // Passwords don't match
                                        return req.StatusConflict(he.NetString("password and password_verification don't match"))
                                }
                        }
                        internalerror := req.StatusBadRequest("Internal Error")
			o.Email = email
			err = o.SetPassword(password); if err != nil { return internalerror }
			err = o.Save()               ; if err != nil { return internalerror }
			return req.StatusOK()

		},
		"DELETE": func(req he.Request) he.Response {
			panic("TODO: API: (*User).Methods()[\"DELETE\"]")
		},
	}
}

// View //////////////////////////////////////////////////////////////

func (o *User) Encoders() map[string]he.Encoder {
	return defaultEncoders(o)
}

// Directory ("Controller") //////////////////////////////////////////

type t_dirUsers struct {
	methods map[string]he.Handler
}

func newDirUsers() t_dirUsers {
	r := t_dirUsers{}
	r.methods = map[string]he.Handler{
		"POST": func(req he.Request) he.Response {
			db := req.Things["db"].(DB)
			badbody := req.StatusBadRequest("submitted body not what expected")
			hash, ok := req.Entity.(map[string]interface{}); if !ok { return badbody }
			username, ok := hash["username"].(string)      ; if !ok { return badbody }
			email   , ok := hash["email"].(string)         ; if !ok { return badbody }
			password, ok := hash["password"].(string)      ; if !ok { return badbody }

			if password2, ok := hash["password_verification"].(string); ok {
				if password != password2 {
					// Passwords don't match
					return req.StatusConflict(he.NetString("password and password_verification don't match"))
				}
			}

			username = strings.ToLower(username)

			user := NewUser(db, username, password, email)
			if user == nil {
				return req.StatusConflict(he.NetString("either that username or password is already taken"))
			} else {
				return req.StatusCreated(r, username)
			}
		},
	}
	return r
}

func (d t_dirUsers) Methods() map[string]he.Handler {
	return d.methods
}

func (d t_dirUsers) Subentity(name string, req he.Request) he.Entity {
	db := req.Things["db"].(DB)
	return GetUserByName(db, name)
}
<|endoftext|>