<commit_before>package rtmp

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"sync"
	//"sync/atomic"
	"time"
)

const (
	DEFAULT_POOL_SIZE = 4096
)

var ErrBufferFull = errors.New("buffer full")

var pool = NewMediaFramePool(DEFAULT_POOL_SIZE)

func NewMediaFrame() *MediaFrame {
	return pool.New()
}

type MediaFrame struct {
	Idx            int
	Timestamp      uint32
	Type           byte //8 audio,9 video
	VideoFrameType byte //4bit
	VideoCodecID   byte //4bit
	AudioFormat    byte //4bit
	SamplingRate   byte //2bit
	SampleLength   byte //1bit
	AudioType      byte //1bit
	Payload        *bytes.Buffer
	StreamId       uint32
	count          int32
	p              *MediaFramePool
}

func (p *MediaFrame) IFrame() bool {
	return p.VideoFrameType == 1 || p.VideoFrameType == 4
}

func (p *MediaFrame) String() string {
	if p == nil {
		return "<nil>"
	}
	if p.Type == RTMP_MSG_AUDIO {
		return fmt.Sprintf("%v Audio Frame Timestamp/%v Type/%v AudioFromat/%v SampleRate/%v SampleLength/%v AudioType/%v Payload/%v StreamId/%v", p.Idx, float64(p.Timestamp)/1000.0, p.Type, audioformat[p.AudioFormat], samplerate[p.SamplingRate], samplelength[p.SampleLength], audiotype[p.AudioType], p.Payload.Len(), p.StreamId)
	} else if p.Type == RTMP_MSG_VIDEO {
		return fmt.Sprintf("%v Video Frame Timestamp/%v Type/%v VideoFrameType/%v VideoCodecID/%v Payload/%v StreamId/%v", p.Idx, float64(p.Timestamp)/1000.0, p.Type, videoframetype[p.VideoFrameType], videocodec[p.VideoCodecID], p.Payload.Len(), p.StreamId)
	}
	return fmt.Sprintf("%v Frame Timestamp/%v Type/%v Payload/%v StreamId/%v", p.Idx, p.Timestamp, p.Type, p.Payload.Len(), p.StreamId)
}

func (o *MediaFrame) Ref() *MediaFrame {
	//atomic.AddInt32(&o.count, 1)
	return o
}

func (o *MediaFrame) Release() {
	// if nc := atomic.AddInt32(&o.count, -1); nc <= 0 {
	// 	select {
	// 	case o.p.pool <- o:
	// 	default:
	// 	}
	// }
}

func (o *MediaFrame) Bytes() []byte {
	return o.Payload.Bytes()
}

func (o *MediaFrame) WriteTo(w io.Writer) (int, error) {
	return w.Write(o.Payload.Bytes())
}

type MediaFramePool struct {
	pool chan *MediaFrame
}

func NewMediaFramePool(size int) *MediaFramePool {
	return &MediaFramePool{pool: make(chan *MediaFrame, size)}
}

func (p *MediaFramePool) New() *MediaFrame {
	var x *MediaFrame
	// select {
	// case x = <-p.pool:
	// 	x.count = 1
	// 	x.Idx = 0
	// 	x.Timestamp = 0
	// 	x.VideoFrameType = 0
	// 	x.VideoCodecID = 0
	// 	x.AudioFormat = 0
	// 	x.SamplingRate = 0
	// 	x.SampleLength = 0
	// 	x.Payload.Reset()
	// default:
	x = &MediaFrame{count: 1, p: p, Payload: bytes.NewBuffer(nil)}
	//}
	return x
}

type NetStream interface {
	NsID() int
	Name() string
	String() string
	Notify(idx *int) error
	Close() error
	StreamObject() *StreamObject
}
type MediaGop struct {
	idx    int
	frames []*MediaFrame
	//freshChunk *RtmpChunker
	//chunk      *RtmpChunker
	//audio bool
	videoConfig *MediaFrame
	audioConfig *MediaFrame
	metaConfig  *MediaFrame
}

func (o *MediaGop) Release() {
	for _, f := range o.frames {
		f.Release()
	}
	o.frames = o.frames[0:0]
	//o.freshChunk.Reset()
	//o.chunk.Reset()
	//o.freshChunk = nil
	//o.chunk = nil
}

func (o *MediaGop) Len() int {
	return len(o.frames)
}

type StreamObject struct {
	name     string
	duration uint32
	list     []int
	cache    map[int]*MediaGop //cmap.ConcurrentMap
	//gopcache map[int]*MediaGop
	subs               []NetStream
	subch              chan NetStream
	sublock            sync.RWMutex
	notify             chan *int
	lock               sync.RWMutex
	idx                int
	gidx               int
	csize              int
	metaData           *MediaFrame
	firstVideoKeyFrame *MediaFrame
	firstAudioKeyFrame *MediaFrame
	//lastVideoKeyFrame  *MediaFrame
	gop      *MediaGop
	streamid uint32
}

func new_streamObject(sid string, timeout time.Duration, record bool, csize int) (obj *StreamObject, err error) {
	obj = &StreamObject{
		name:   sid,
		list:   []int{},
		cache:  make(map[int]*MediaGop, csize),
		subs:   []NetStream{},
		notify: make(chan *int, csize*100),
		csize:  csize,
	}
	addObject(obj)
	go obj.loop(timeout)
	return obj, nil
}

func (m *StreamObject) Attach(c NetStream) {
	m.sublock.Lock()
	m.subs = append(m.subs, c)
	m.sublock.Unlock()
}
func (m *StreamObject) ReadGop(idx *int) *MediaGop {
	m.lock.RLock()
	if s, found := m.cache[*idx]; found {
		m.lock.RUnlock()
		return s
	}
	m.lock.RUnlock()
	log.Warn("Gop", m.name, *idx, "Not Found")
	return nil
}

func (m *StreamObject) WriteFrame(s *MediaFrame) (err error) {
	m.lock.Lock()
	if m.idx >= 0xffffffffffffff {
		m.idx = 0
	}
	s.Idx = m.idx
	m.idx += 1
	m.duration = s.Timestamp
	if s.Type == RTMP_MSG_VIDEO && s.IFrame() && m.firstVideoKeyFrame == nil {
		log.Info(">>>>", s)
		m.firstVideoKeyFrame = s
		m.streamid = s.StreamId
		m.lock.Unlock()
		return
	}
	if s.Type == RTMP_MSG_AUDIO && m.firstAudioKeyFrame == nil {
		log.Info(">>>>", s)
		m.firstAudioKeyFrame = s
		m.lock.Unlock()
		return
	}
	if s.Type == RTMP_MSG_AMF_META && m.metaData == nil {
		log.Info(">>>>", s)
		m.metaData = s
		m.lock.Unlock()
		return
	}
	if m.gop == nil {
		m.gop = &MediaGop{0, make([]*MediaFrame, 0), m.firstVideoKeyFrame, m.firstAudioKeyFrame, m.metaData}
	}

	if len(m.list) >= m.csize {
		idx := m.list[0]
		if s, found := m.cache[idx]; found {
			s.Release()
			delete(m.cache, idx)
		}
		m.list = m.list[1:]
	}
	if s.IFrame() && m.gop.Len() > 0 {
		gop := m.gop
		m.list = append(m.list, gop.idx)
		m.cache[gop.idx] = gop
		log.Info("Gop", m.name, gop.idx, gop.Len(), len(m.list))
		m.gop = &MediaGop{gop.idx + 1, []*MediaFrame{s}, m.firstVideoKeyFrame, m.firstAudioKeyFrame, m.metaData}
		// m.gop.chunk.wchunks = gop.chunk.wchunks
		// m.gop.freshChunk.writeMetadata(m.metaData)
		// m.gop.freshChunk.writeFullVideo(m.firstVideoKeyFrame)
		// m.gop.freshChunk.writeFullAudio(m.firstAudioKeyFrame)
		// m.gop.freshChunk.writeFullVideo(s)
		// m.gop.chunk.writeVideo(s)
		m.lock.Unlock()
		select {
		case m.notify <- &gop.idx:
		default:
			err = ErrBufferFull
		}
		return
	}
	m.gop.frames = append(m.gop.frames, s)
	// if s.Type == RTMP_MSG_VIDEO {
	// 	m.gop.freshChunk.writeVideo(s)
	// 	m.gop.chunk.writeVideo(s)
	// } else if s.Type == RTMP_MSG_AUDIO {
	// 	if !m.gop.audio {
	// 		m.gop.freshChunk.writeFullAudio(s)
	// 	} else {
	// 		m.gop.freshChunk.writeAudio(s)
	// 	}
	// 	m.gop.chunk.writeAudio(s)
	// }
	m.lock.Unlock()
	return
}

func (m *StreamObject) Close() {
	removeObject(m.name)
	close(m.notify)
}
func (m *StreamObject) loop(timeout time.Duration) {
	log.Info(m.name, "stream object is runing")
	defer log.Info(m.name, "stream object is stopped")
	var (
		opened bool
		idx    *int
		w      NetStream
		err    error
		nsubs  = []NetStream{}
		subs   = []NetStream{}
	)
	defer m.clear()
	for {
		select {
		case idx, opened = <-m.notify:
			if !opened {
				return
			}
			m.sublock.Lock()
			nsubs = nsubs[0:0]
			subs = m.subs[:]
			m.sublock.Unlock()
			log.Info("players", m.name, len(subs))
			for _, w = range subs {
				if err = w.Notify(idx); err != nil {
					log.Error(w, err)
					w.Close()
				} else {
					nsubs = append(nsubs, w)
				}
			}
			m.sublock.Lock()
			m.subs = nsubs[:]
			m.sublock.Unlock()
		case <-time.After(timeout):
			m.Close()
		}
	}
}

func (m *StreamObject) clear() {
	m.sublock.Lock()
	for _, w := range m.subs {
		w.Close()
	}
	m.subs = m.subs[0:0]
	m.sublock.Unlock()
}
<commit_msg>fix int overflow problem<commit_after>package rtmp

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"sync"
	//"sync/atomic"
	"time"
)

const (
	DEFAULT_POOL_SIZE = 4096
)

var ErrBufferFull = errors.New("buffer full")

var pool = NewMediaFramePool(DEFAULT_POOL_SIZE)

func NewMediaFrame() *MediaFrame {
	return pool.New()
}

type MediaFrame struct {
	Idx            int
	Timestamp      uint32
	Type           byte //8 audio,9 video
	VideoFrameType byte //4bit
	VideoCodecID   byte //4bit
	AudioFormat    byte //4bit
	SamplingRate   byte //2bit
	SampleLength   byte //1bit
	AudioType      byte //1bit
	Payload        *bytes.Buffer
	StreamId       uint32
	count          int32
	p              *MediaFramePool
}

func (p *MediaFrame) IFrame() bool {
	return p.VideoFrameType == 1 || p.VideoFrameType == 4
}

func (p *MediaFrame) String() string {
	if p == nil {
		return "<nil>"
	}
	if p.Type == RTMP_MSG_AUDIO {
		return fmt.Sprintf("%v Audio Frame Timestamp/%v Type/%v AudioFromat/%v SampleRate/%v SampleLength/%v AudioType/%v Payload/%v StreamId/%v", p.Idx, float64(p.Timestamp)/1000.0, p.Type, audioformat[p.AudioFormat], samplerate[p.SamplingRate], samplelength[p.SampleLength], audiotype[p.AudioType], p.Payload.Len(), p.StreamId)
	} else if p.Type == RTMP_MSG_VIDEO {
		return fmt.Sprintf("%v Video Frame Timestamp/%v Type/%v VideoFrameType/%v VideoCodecID/%v Payload/%v StreamId/%v", p.Idx, float64(p.Timestamp)/1000.0, p.Type, videoframetype[p.VideoFrameType], videocodec[p.VideoCodecID], p.Payload.Len(), p.StreamId)
	}
	return fmt.Sprintf("%v Frame Timestamp/%v Type/%v Payload/%v StreamId/%v", p.Idx, p.Timestamp, p.Type, p.Payload.Len(), p.StreamId)
}

func (o *MediaFrame) Ref() *MediaFrame {
	//atomic.AddInt32(&o.count, 1)
	return o
}

func (o *MediaFrame) Release() {
	// if nc := atomic.AddInt32(&o.count, -1); nc <= 0 {
	// 	select {
	// 	case o.p.pool <- o:
	// 	default:
	// 	}
	// }
}

func (o *MediaFrame) Bytes() []byte {
	return o.Payload.Bytes()
}

func (o *MediaFrame) WriteTo(w io.Writer) (int, error) {
	return w.Write(o.Payload.Bytes())
}

type MediaFramePool struct {
	pool chan *MediaFrame
}

func NewMediaFramePool(size int) *MediaFramePool {
	return &MediaFramePool{pool: make(chan *MediaFrame, size)}
}

func (p *MediaFramePool) New() *MediaFrame {
	var x *MediaFrame
	// select {
	// case x = <-p.pool:
	// 	x.count = 1
	// 	x.Idx = 0
	// 	x.Timestamp = 0
	// 	x.VideoFrameType = 0
	// 	x.VideoCodecID = 0
	// 	x.AudioFormat = 0
	// 	x.SamplingRate = 0
	// 	x.SampleLength = 0
	// 	x.Payload.Reset()
	// default:
	x = &MediaFrame{count: 1, p: p, Payload: bytes.NewBuffer(nil)}
	//}
	return x
}

type NetStream interface {
	NsID() int
	Name() string
	String() string
	Notify(idx *int) error
	Close() error
	StreamObject() *StreamObject
}
type MediaGop struct {
	idx    int
	frames []*MediaFrame
	//freshChunk *RtmpChunker
	//chunk      *RtmpChunker
	//audio bool
	videoConfig *MediaFrame
	audioConfig *MediaFrame
	metaConfig  *MediaFrame
}

func (o *MediaGop) Release() {
	for _, f := range o.frames {
		f.Release()
	}
	o.frames = o.frames[0:0]
	//o.freshChunk.Reset()
	//o.chunk.Reset()
	//o.freshChunk = nil
	//o.chunk = nil
}

func (o *MediaGop) Len() int {
	return len(o.frames)
}

type StreamObject struct {
	name     string
	duration uint32
	list     []int
	cache    map[int]*MediaGop //cmap.ConcurrentMap
	//gopcache map[int]*MediaGop
	subs               []NetStream
	subch              chan NetStream
	sublock            sync.RWMutex
	notify             chan *int
	lock               sync.RWMutex
	idx                int
	gidx               int
	csize              int
	metaData           *MediaFrame
	firstVideoKeyFrame *MediaFrame
	firstAudioKeyFrame *MediaFrame
	//lastVideoKeyFrame  *MediaFrame
	gop      *MediaGop
	streamid uint32
}

func new_streamObject(sid string, timeout time.Duration, record bool, csize int) (obj *StreamObject, err error) {
	obj = &StreamObject{
		name:   sid,
		list:   []int{},
		cache:  make(map[int]*MediaGop, csize),
		subs:   []NetStream{},
		notify: make(chan *int, csize*100),
		csize:  csize,
	}
	addObject(obj)
	go obj.loop(timeout)
	return obj, nil
}

func (m *StreamObject) Attach(c NetStream) {
	m.sublock.Lock()
	m.subs = append(m.subs, c)
	m.sublock.Unlock()
}
func (m *StreamObject) ReadGop(idx *int) *MediaGop {
	m.lock.RLock()
	if s, found := m.cache[*idx]; found {
		m.lock.RUnlock()
		return s
	}
	m.lock.RUnlock()
	log.Warn("Gop", m.name, *idx, "Not Found")
	return nil
}

func (m *StreamObject) WriteFrame(s *MediaFrame) (err error) {
	m.lock.Lock()
	if m.idx >= 0x7fffffff {
		m.idx = 0
	}
	s.Idx = m.idx
	m.idx += 1
	m.duration = s.Timestamp
	if s.Type == RTMP_MSG_VIDEO && s.IFrame() && m.firstVideoKeyFrame == nil {
		log.Info(">>>>", s)
		m.firstVideoKeyFrame = s
		m.streamid = s.StreamId
		m.lock.Unlock()
		return
	}
	if s.Type == RTMP_MSG_AUDIO && m.firstAudioKeyFrame == nil {
		log.Info(">>>>", s)
		m.firstAudioKeyFrame = s
		m.lock.Unlock()
		return
	}
	if s.Type == RTMP_MSG_AMF_META && m.metaData == nil {
		log.Info(">>>>", s)
		m.metaData = s
		m.lock.Unlock()
		return
	}
	if m.gop == nil {
		m.gop = &MediaGop{0, make([]*MediaFrame, 0), m.firstVideoKeyFrame, m.firstAudioKeyFrame, m.metaData}
	}

	if len(m.list) >= m.csize {
		idx := m.list[0]
		if s, found := m.cache[idx]; found {
			s.Release()
			delete(m.cache, idx)
		}
		m.list = m.list[1:]
	}
	if s.IFrame() && m.gop.Len() > 0 {
		gop := m.gop
		m.list = append(m.list, gop.idx)
		m.cache[gop.idx] = gop
		log.Info("Gop", m.name, gop.idx, gop.Len(), len(m.list))
		m.gop = &MediaGop{gop.idx + 1, []*MediaFrame{s}, m.firstVideoKeyFrame, m.firstAudioKeyFrame, m.metaData}
		// m.gop.chunk.wchunks = gop.chunk.wchunks
		// m.gop.freshChunk.writeMetadata(m.metaData)
		// m.gop.freshChunk.writeFullVideo(m.firstVideoKeyFrame)
		// m.gop.freshChunk.writeFullAudio(m.firstAudioKeyFrame)
		// m.gop.freshChunk.writeFullVideo(s)
		// m.gop.chunk.writeVideo(s)
		m.lock.Unlock()
		select {
		case m.notify <- &gop.idx:
		default:
			err = ErrBufferFull
		}
		return
	}
	m.gop.frames = append(m.gop.frames, s)
	// if s.Type == RTMP_MSG_VIDEO {
	// 	m.gop.freshChunk.writeVideo(s)
	// 	m.gop.chunk.writeVideo(s)
	// } else if s.Type == RTMP_MSG_AUDIO {
	// 	if !m.gop.audio {
	// 		m.gop.freshChunk.writeFullAudio(s)
	// 	} else {
	// 		m.gop.freshChunk.writeAudio(s)
	// 	}
	// 	m.gop.chunk.writeAudio(s)
	// }
	m.lock.Unlock()
	return
}

func (m *StreamObject) Close() {
	removeObject(m.name)
	close(m.notify)
}
func (m *StreamObject) loop(timeout time.Duration) {
	log.Info(m.name, "stream object is runing")
	defer log.Info(m.name, "stream object is stopped")
	var (
		opened bool
		idx    *int
		w      NetStream
		err    error
		nsubs  = []NetStream{}
		subs   = []NetStream{}
	)
	defer m.clear()
	for {
		select {
		case idx, opened = <-m.notify:
			if !opened {
				return
			}
			m.sublock.Lock()
			nsubs = nsubs[0:0]
			subs = m.subs[:]
			m.sublock.Unlock()
			log.Info("players", m.name, len(subs))
			for _, w = range subs {
				if err = w.Notify(idx); err != nil {
					log.Error(w, err)
					w.Close()
				} else {
					nsubs = append(nsubs, w)
				}
			}
			m.sublock.Lock()
			m.subs = nsubs[:]
			m.sublock.Unlock()
		case <-time.After(timeout):
			m.Close()
		}
	}
}

func (m *StreamObject) clear() {
	m.sublock.Lock()
	for _, w := range m.subs {
		w.Close()
	}
	m.subs = m.subs[0:0]
	m.sublock.Unlock()
}
<|endoftext|>


<commit_before>package s3api

import (
	"crypto/md5"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"sort"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/s3api/s3err"

	"github.com/gorilla/mux"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	weed_server "github.com/chrislusf/seaweedfs/weed/server"
	"github.com/chrislusf/seaweedfs/weed/util"
)

var (
	client *http.Client
)

func init() {
	client = &http.Client{Transport: &http.Transport{
		MaxIdleConns:        1024,
		MaxIdleConnsPerHost: 1024,
	}}
}

func (s3a *S3ApiServer) PutObjectHandler(w http.ResponseWriter, r *http.Request) {

	// http://docs.aws.amazon.com/AmazonS3/latest/dev/UploadingObjects.html

	bucket, object := getBucketAndObject(r)

	_, err := validateContentMd5(r.Header)
	if err != nil {
		writeErrorResponse(w, s3err.ErrInvalidDigest, r.URL)
		return
	}

	dataReader := r.Body
	if s3a.iam.isEnabled() {
		rAuthType := getRequestAuthType(r)
		var s3ErrCode s3err.ErrorCode
		switch rAuthType {
		case authTypeStreamingSigned:
			dataReader, s3ErrCode = s3a.iam.newSignV4ChunkedReader(r)
		case authTypeSignedV2, authTypePresignedV2:
			_, s3ErrCode = s3a.iam.isReqAuthenticatedV2(r)
		case authTypePresigned, authTypeSigned:
			_, s3ErrCode = s3a.iam.reqSignatureV4Verify(r)
		}
		if s3ErrCode != s3err.ErrNone {
			writeErrorResponse(w, s3ErrCode, r.URL)
			return
		}
	}
	defer dataReader.Close()

	if strings.HasSuffix(object, "/") {
		if err := s3a.mkdir(s3a.option.BucketsPath, bucket+object, nil); err != nil {
			writeErrorResponse(w, s3err.ErrInternalError, r.URL)
			return
		}
	} else {
		uploadUrl := fmt.Sprintf("http://%s%s/%s%s", s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

		etag, errCode := s3a.putToFiler(r, uploadUrl, dataReader)

		if errCode != s3err.ErrNone {
			writeErrorResponse(w, errCode, r.URL)
			return
		}

		setEtag(w, etag)
	}

	writeSuccessResponseEmpty(w)
}

func (s3a *S3ApiServer) GetObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	if strings.HasSuffix(r.URL.Path, "/") {
		writeErrorResponse(w, s3err.ErrNotImplemented, r.URL)
		return
	}

	destUrl := fmt.Sprintf("http://%s%s/%s%s",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

	s3a.proxyToFiler(w, r, destUrl, passThroughResponse)

}

func (s3a *S3ApiServer) HeadObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	destUrl := fmt.Sprintf("http://%s%s/%s%s",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

	s3a.proxyToFiler(w, r, destUrl, passThroughResponse)

}

func (s3a *S3ApiServer) DeleteObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	destUrl := fmt.Sprintf("http://%s%s/%s%s?recursive=true",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

	s3a.proxyToFiler(w, r, destUrl, func(proxyResponse *http.Response, w http.ResponseWriter) {
		for k, v := range proxyResponse.Header {
			w.Header()[k] = v
		}
		w.WriteHeader(http.StatusNoContent)
	})
}

// / ObjectIdentifier carries key name for the object to delete.
type ObjectIdentifier struct {
	ObjectName string `xml:"Key"`
}

// DeleteObjectsRequest - xml carrying the object key names which needs to be deleted.
type DeleteObjectsRequest struct {
	// Element to enable quiet mode for the request
	Quiet bool
	// List of objects to be deleted
	Objects []ObjectIdentifier `xml:"Object"`
}

// DeleteError structure.
type DeleteError struct {
	Code    string
	Message string
	Key     string
}

// DeleteObjectsResponse container for multiple object deletes.
type DeleteObjectsResponse struct {
	XMLName xml.Name `xml:"http://s3.amazonaws.com/doc/2006-03-01/ DeleteResult" json:"-"`

	// Collection of all deleted objects
	DeletedObjects []ObjectIdentifier `xml:"Deleted,omitempty"`

	// Collection of errors deleting certain objects.
	Errors []DeleteError `xml:"Error,omitempty"`
}

// DeleteMultipleObjectsHandler - Delete multiple objects
func (s3a *S3ApiServer) DeleteMultipleObjectsHandler(w http.ResponseWriter, r *http.Request) {

	bucket, _ := getBucketAndObject(r)

	deleteXMLBytes, err := ioutil.ReadAll(r.Body)
	if err != nil {
		writeErrorResponse(w, s3err.ErrInternalError, r.URL)
		return
	}

	deleteObjects := &DeleteObjectsRequest{}
	if err := xml.Unmarshal(deleteXMLBytes, deleteObjects); err != nil {
		writeErrorResponse(w, s3err.ErrMalformedXML, r.URL)
		return
	}

	var deletedObjects []ObjectIdentifier
	var deleteErrors []DeleteError

	directoriesWithDeletion := make(map[string]int)

	s3a.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		// delete file entries
		for _, object := range deleteObjects.Objects {

			lastSeparator := strings.LastIndex(object.ObjectName, "/")
			parentDirectoryPath, entryName, isDeleteData, isRecursive := "/", object.ObjectName, true, false
			if lastSeparator > 0 && lastSeparator+1 < len(object.ObjectName) {
				entryName = object.ObjectName[lastSeparator+1:]
				parentDirectoryPath = "/" + object.ObjectName[:lastSeparator]
			}
			parentDirectoryPath = fmt.Sprintf("%s/%s%s", s3a.option.BucketsPath, bucket, parentDirectoryPath)

			err := doDeleteEntry(client, parentDirectoryPath, entryName, isDeleteData, isRecursive)
			if err == nil {
				directoriesWithDeletion[parentDirectoryPath]++
				deletedObjects = append(deletedObjects, object)
			} else {
				delete(directoriesWithDeletion, parentDirectoryPath)
				deleteErrors = append(deleteErrors, DeleteError{
					Code:    "",
					Message: err.Error(),
					Key:     object.ObjectName,
				})
			}
		}

		// purge empty folders, only checking folders with deletions
		for len(directoriesWithDeletion) > 0 {
			directoriesWithDeletion = doDeleteEmptyDirectories(client, directoriesWithDeletion)
		}

		return nil
	})

	deleteResp := DeleteObjectsResponse{}
	if !deleteObjects.Quiet {
		deleteResp.DeletedObjects = deletedObjects
	}
	deleteResp.Errors = deleteErrors

	writeSuccessResponseXML(w, encodeResponse(deleteResp))

}

func doDeleteEmptyDirectories(client filer_pb.SeaweedFilerClient, directoriesWithDeletion map[string]int) (newDirectoriesWithDeletion map[string]int) {
	var allDirs []string
	for dir, _ := range directoriesWithDeletion {
		allDirs = append(allDirs, dir)
	}
	sort.Slice(allDirs, func(i, j int) bool {
		return len(allDirs[i]) > len(allDirs[j])
	})
	newDirectoriesWithDeletion = make(map[string]int)
	for _, dir := range allDirs {
		parentDir, dirName := util.FullPath(dir).DirAndName()
		if err := doDeleteEntry(client, parentDir, dirName, false, false); err != nil {
			glog.V(4).Infof("directory %s has %d deletion but still not empty: %v", dir, directoriesWithDeletion[dir], err)
		} else {
			newDirectoriesWithDeletion[parentDir]++
		}
	}
	return
}

var passThroughHeaders = []string{
	"response-cache-control",
	"response-content-disposition",
	"response-content-encoding",
	"response-content-language",
	"response-content-type",
	"response-expires",
}

func (s3a *S3ApiServer) proxyToFiler(w http.ResponseWriter, r *http.Request, destUrl string, responseFn func(proxyResponse *http.Response, w http.ResponseWriter)) {

	glog.V(2).Infof("s3 proxying %s to %s", r.Method, destUrl)

	proxyReq, err := http.NewRequest(r.Method, destUrl, r.Body)

	if err != nil {
		glog.Errorf("NewRequest %s: %v", destUrl, err)
		writeErrorResponse(w, s3err.ErrInternalError, r.URL)
		return
	}

	proxyReq.Header.Set("Host", s3a.option.Filer)
	proxyReq.Header.Set("X-Forwarded-For", r.RemoteAddr)

	for header, values := range r.Header {
		// handle s3 related headers
		passed := false
		for _, h := range passThroughHeaders {
			if strings.ToLower(header) == h && len(values) > 0 {
				proxyReq.Header.Add(header[len("response-"):], values[0])
				passed = true
				break
			}
		}
		if passed {
			continue
		}
		// handle other headers
		for _, value := range values {
			proxyReq.Header.Add(header, value)
		}
	}

	resp, postErr := client.Do(proxyReq)

	if postErr != nil {
		glog.Errorf("post to filer: %v", postErr)
		writeErrorResponse(w, s3err.ErrInternalError, r.URL)
		return
	}
	defer util.CloseResponse(resp)

	if (resp.ContentLength == -1 || resp.StatusCode == 404) && !strings.HasSuffix(destUrl, "/") {
		if r.Method != "DELETE" {
			writeErrorResponse(w, s3err.ErrNoSuchKey, r.URL)
			return
		}
	}

	responseFn(resp, w)

}

func passThroughResponse(proxyResponse *http.Response, w http.ResponseWriter) {
	for k, v := range proxyResponse.Header {
		w.Header()[k] = v
	}
	w.WriteHeader(proxyResponse.StatusCode)
	io.Copy(w, proxyResponse.Body)
}

func (s3a *S3ApiServer) putToFiler(r *http.Request, uploadUrl string, dataReader io.Reader) (etag string, code s3err.ErrorCode) {

	hash := md5.New()
	var body = io.TeeReader(dataReader, hash)

	proxyReq, err := http.NewRequest("PUT", uploadUrl, body)

	if err != nil {
		glog.Errorf("NewRequest %s: %v", uploadUrl, err)
		return "", s3err.ErrInternalError
	}

	proxyReq.Header.Set("Host", s3a.option.Filer)
	proxyReq.Header.Set("X-Forwarded-For", r.RemoteAddr)

	for header, values := range r.Header {
		for _, value := range values {
			proxyReq.Header.Add(header, value)
		}
	}

	resp, postErr := client.Do(proxyReq)

	if postErr != nil {
		glog.Errorf("post to filer: %v", postErr)
		return "", s3err.ErrInternalError
	}
	defer resp.Body.Close()

	etag = fmt.Sprintf("%x", hash.Sum(nil))

	resp_body, ra_err := ioutil.ReadAll(resp.Body)
	if ra_err != nil {
		glog.Errorf("upload to filer response read %d: %v", resp.StatusCode, ra_err)
		return etag, s3err.ErrInternalError
	}
	var ret weed_server.FilerPostResult
	unmarshal_err := json.Unmarshal(resp_body, &ret)
	if unmarshal_err != nil {
		glog.Errorf("failing to read upload to %s : %v", uploadUrl, string(resp_body))
		return "", s3err.ErrInternalError
	}
	if ret.Error != "" {
		glog.Errorf("upload to filer error: %v", ret.Error)
		return "", filerErrorToS3Error(ret.Error)
	}

	return etag, s3err.ErrNone
}

func setEtag(w http.ResponseWriter, etag string) {
	if etag != "" {
		if strings.HasPrefix(etag, "\"") {
			w.Header().Set("ETag", etag)
		} else {
			w.Header().Set("ETag", "\""+etag+"\"")
		}
	}
}

func getBucketAndObject(r *http.Request) (bucket, object string) {
	vars := mux.Vars(r)
	bucket = vars["bucket"]
	object = vars["object"]
	if !strings.HasPrefix(object, "/") {
		object = "/" + object
	}

	return
}

func filerErrorToS3Error(errString string) s3err.ErrorCode {
	if strings.HasPrefix(errString, "existing ") && strings.HasSuffix(errString, "is a directory") {
		return s3err.ErrExistingObjectIsDirectory
	}
	return s3err.ErrInternalError
}
<commit_msg>s3: Critical - S3 Delete Multiple Objects, delete the bucket<commit_after>package s3api

import (
	"crypto/md5"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"sort"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/s3api/s3err"

	"github.com/gorilla/mux"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	weed_server "github.com/chrislusf/seaweedfs/weed/server"
	"github.com/chrislusf/seaweedfs/weed/util"
)

var (
	client *http.Client
)

func init() {
	client = &http.Client{Transport: &http.Transport{
		MaxIdleConns:        1024,
		MaxIdleConnsPerHost: 1024,
	}}
}

func (s3a *S3ApiServer) PutObjectHandler(w http.ResponseWriter, r *http.Request) {

	// http://docs.aws.amazon.com/AmazonS3/latest/dev/UploadingObjects.html

	bucket, object := getBucketAndObject(r)

	_, err := validateContentMd5(r.Header)
	if err != nil {
		writeErrorResponse(w, s3err.ErrInvalidDigest, r.URL)
		return
	}

	dataReader := r.Body
	if s3a.iam.isEnabled() {
		rAuthType := getRequestAuthType(r)
		var s3ErrCode s3err.ErrorCode
		switch rAuthType {
		case authTypeStreamingSigned:
			dataReader, s3ErrCode = s3a.iam.newSignV4ChunkedReader(r)
		case authTypeSignedV2, authTypePresignedV2:
			_, s3ErrCode = s3a.iam.isReqAuthenticatedV2(r)
		case authTypePresigned, authTypeSigned:
			_, s3ErrCode = s3a.iam.reqSignatureV4Verify(r)
		}
		if s3ErrCode != s3err.ErrNone {
			writeErrorResponse(w, s3ErrCode, r.URL)
			return
		}
	}
	defer dataReader.Close()

	if strings.HasSuffix(object, "/") {
		if err := s3a.mkdir(s3a.option.BucketsPath, bucket+object, nil); err != nil {
			writeErrorResponse(w, s3err.ErrInternalError, r.URL)
			return
		}
	} else {
		uploadUrl := fmt.Sprintf("http://%s%s/%s%s", s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

		etag, errCode := s3a.putToFiler(r, uploadUrl, dataReader)

		if errCode != s3err.ErrNone {
			writeErrorResponse(w, errCode, r.URL)
			return
		}

		setEtag(w, etag)
	}

	writeSuccessResponseEmpty(w)
}

func (s3a *S3ApiServer) GetObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	if strings.HasSuffix(r.URL.Path, "/") {
		writeErrorResponse(w, s3err.ErrNotImplemented, r.URL)
		return
	}

	destUrl := fmt.Sprintf("http://%s%s/%s%s",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

	s3a.proxyToFiler(w, r, destUrl, passThroughResponse)

}

func (s3a *S3ApiServer) HeadObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	destUrl := fmt.Sprintf("http://%s%s/%s%s",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

	s3a.proxyToFiler(w, r, destUrl, passThroughResponse)

}

func (s3a *S3ApiServer) DeleteObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	destUrl := fmt.Sprintf("http://%s%s/%s%s?recursive=true",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

	s3a.proxyToFiler(w, r, destUrl, func(proxyResponse *http.Response, w http.ResponseWriter) {
		for k, v := range proxyResponse.Header {
			w.Header()[k] = v
		}
		w.WriteHeader(http.StatusNoContent)
	})
}

// / ObjectIdentifier carries key name for the object to delete.
type ObjectIdentifier struct {
	ObjectName string `xml:"Key"`
}

// DeleteObjectsRequest - xml carrying the object key names which needs to be deleted.
type DeleteObjectsRequest struct {
	// Element to enable quiet mode for the request
	Quiet bool
	// List of objects to be deleted
	Objects []ObjectIdentifier `xml:"Object"`
}

// DeleteError structure.
type DeleteError struct {
	Code    string
	Message string
	Key     string
}

// DeleteObjectsResponse container for multiple object deletes.
type DeleteObjectsResponse struct {
	XMLName xml.Name `xml:"http://s3.amazonaws.com/doc/2006-03-01/ DeleteResult" json:"-"`

	// Collection of all deleted objects
	DeletedObjects []ObjectIdentifier `xml:"Deleted,omitempty"`

	// Collection of errors deleting certain objects.
	Errors []DeleteError `xml:"Error,omitempty"`
}

// DeleteMultipleObjectsHandler - Delete multiple objects
func (s3a *S3ApiServer) DeleteMultipleObjectsHandler(w http.ResponseWriter, r *http.Request) {

	bucket, _ := getBucketAndObject(r)

	deleteXMLBytes, err := ioutil.ReadAll(r.Body)
	if err != nil {
		writeErrorResponse(w, s3err.ErrInternalError, r.URL)
		return
	}

	deleteObjects := &DeleteObjectsRequest{}
	if err := xml.Unmarshal(deleteXMLBytes, deleteObjects); err != nil {
		writeErrorResponse(w, s3err.ErrMalformedXML, r.URL)
		return
	}

	var deletedObjects []ObjectIdentifier
	var deleteErrors []DeleteError

	directoriesWithDeletion := make(map[string]int)

	s3a.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		// delete file entries
		for _, object := range deleteObjects.Objects {

			lastSeparator := strings.LastIndex(object.ObjectName, "/")
			parentDirectoryPath, entryName, isDeleteData, isRecursive := "", object.ObjectName, true, false
			if lastSeparator > 0 && lastSeparator+1 < len(object.ObjectName) {
				entryName = object.ObjectName[lastSeparator+1:]
				parentDirectoryPath = "/" + object.ObjectName[:lastSeparator]
			}
			parentDirectoryPath = fmt.Sprintf("%s/%s%s", s3a.option.BucketsPath, bucket, parentDirectoryPath)

			err := doDeleteEntry(client, parentDirectoryPath, entryName, isDeleteData, isRecursive)
			if err == nil {
				directoriesWithDeletion[parentDirectoryPath]++
				deletedObjects = append(deletedObjects, object)
			} else {
				delete(directoriesWithDeletion, parentDirectoryPath)
				deleteErrors = append(deleteErrors, DeleteError{
					Code:    "",
					Message: err.Error(),
					Key:     object.ObjectName,
				})
			}
		}

		// purge empty folders, only checking folders with deletions
		for len(directoriesWithDeletion) > 0 {
			directoriesWithDeletion = s3a.doDeleteEmptyDirectories(client, directoriesWithDeletion)
		}

		return nil
	})

	deleteResp := DeleteObjectsResponse{}
	if !deleteObjects.Quiet {
		deleteResp.DeletedObjects = deletedObjects
	}
	deleteResp.Errors = deleteErrors

	writeSuccessResponseXML(w, encodeResponse(deleteResp))

}

func (s3a *S3ApiServer) doDeleteEmptyDirectories(client filer_pb.SeaweedFilerClient, directoriesWithDeletion map[string]int) (newDirectoriesWithDeletion map[string]int) {
	var allDirs []string
	for dir, _ := range directoriesWithDeletion {
		allDirs = append(allDirs, dir)
	}
	sort.Slice(allDirs, func(i, j int) bool {
		return len(allDirs[i]) > len(allDirs[j])
	})
	newDirectoriesWithDeletion = make(map[string]int)
	for _, dir := range allDirs {
		parentDir, dirName := util.FullPath(dir).DirAndName()
		if parentDir == s3a.option.BucketsPath {
			continue
		}
		if err := doDeleteEntry(client, parentDir, dirName, false, false); err != nil {
			glog.V(4).Infof("directory %s has %d deletion but still not empty: %v", dir, directoriesWithDeletion[dir], err)
		} else {
			newDirectoriesWithDeletion[parentDir]++
		}
	}
	return
}

var passThroughHeaders = []string{
	"response-cache-control",
	"response-content-disposition",
	"response-content-encoding",
	"response-content-language",
	"response-content-type",
	"response-expires",
}

func (s3a *S3ApiServer) proxyToFiler(w http.ResponseWriter, r *http.Request, destUrl string, responseFn func(proxyResponse *http.Response, w http.ResponseWriter)) {

	glog.V(2).Infof("s3 proxying %s to %s", r.Method, destUrl)

	proxyReq, err := http.NewRequest(r.Method, destUrl, r.Body)

	if err != nil {
		glog.Errorf("NewRequest %s: %v", destUrl, err)
		writeErrorResponse(w, s3err.ErrInternalError, r.URL)
		return
	}

	proxyReq.Header.Set("Host", s3a.option.Filer)
	proxyReq.Header.Set("X-Forwarded-For", r.RemoteAddr)

	for header, values := range r.Header {
		// handle s3 related headers
		passed := false
		for _, h := range passThroughHeaders {
			if strings.ToLower(header) == h && len(values) > 0 {
				proxyReq.Header.Add(header[len("response-"):], values[0])
				passed = true
				break
			}
		}
		if passed {
			continue
		}
		// handle other headers
		for _, value := range values {
			proxyReq.Header.Add(header, value)
		}
	}

	resp, postErr := client.Do(proxyReq)

	if postErr != nil {
		glog.Errorf("post to filer: %v", postErr)
		writeErrorResponse(w, s3err.ErrInternalError, r.URL)
		return
	}
	defer util.CloseResponse(resp)

	if (resp.ContentLength == -1 || resp.StatusCode == 404) && !strings.HasSuffix(destUrl, "/") {
		if r.Method != "DELETE" {
			writeErrorResponse(w, s3err.ErrNoSuchKey, r.URL)
			return
		}
	}

	responseFn(resp, w)

}

func passThroughResponse(proxyResponse *http.Response, w http.ResponseWriter) {
	for k, v := range proxyResponse.Header {
		w.Header()[k] = v
	}
	w.WriteHeader(proxyResponse.StatusCode)
	io.Copy(w, proxyResponse.Body)
}

func (s3a *S3ApiServer) putToFiler(r *http.Request, uploadUrl string, dataReader io.Reader) (etag string, code s3err.ErrorCode) {

	hash := md5.New()
	var body = io.TeeReader(dataReader, hash)

	proxyReq, err := http.NewRequest("PUT", uploadUrl, body)

	if err != nil {
		glog.Errorf("NewRequest %s: %v", uploadUrl, err)
		return "", s3err.ErrInternalError
	}

	proxyReq.Header.Set("Host", s3a.option.Filer)
	proxyReq.Header.Set("X-Forwarded-For", r.RemoteAddr)

	for header, values := range r.Header {
		for _, value := range values {
			proxyReq.Header.Add(header, value)
		}
	}

	resp, postErr := client.Do(proxyReq)

	if postErr != nil {
		glog.Errorf("post to filer: %v", postErr)
		return "", s3err.ErrInternalError
	}
	defer resp.Body.Close()

	etag = fmt.Sprintf("%x", hash.Sum(nil))

	resp_body, ra_err := ioutil.ReadAll(resp.Body)
	if ra_err != nil {
		glog.Errorf("upload to filer response read %d: %v", resp.StatusCode, ra_err)
		return etag, s3err.ErrInternalError
	}
	var ret weed_server.FilerPostResult
	unmarshal_err := json.Unmarshal(resp_body, &ret)
	if unmarshal_err != nil {
		glog.Errorf("failing to read upload to %s : %v", uploadUrl, string(resp_body))
		return "", s3err.ErrInternalError
	}
	if ret.Error != "" {
		glog.Errorf("upload to filer error: %v", ret.Error)
		return "", filerErrorToS3Error(ret.Error)
	}

	return etag, s3err.ErrNone
}

func setEtag(w http.ResponseWriter, etag string) {
	if etag != "" {
		if strings.HasPrefix(etag, "\"") {
			w.Header().Set("ETag", etag)
		} else {
			w.Header().Set("ETag", "\""+etag+"\"")
		}
	}
}

func getBucketAndObject(r *http.Request) (bucket, object string) {
	vars := mux.Vars(r)
	bucket = vars["bucket"]
	object = vars["object"]
	if !strings.HasPrefix(object, "/") {
		object = "/" + object
	}

	return
}

func filerErrorToS3Error(errString string) s3err.ErrorCode {
	if strings.HasPrefix(errString, "existing ") && strings.HasSuffix(errString, "is a directory") {
		return s3err.ErrExistingObjectIsDirectory
	}
	return s3err.ErrInternalError
}
<|endoftext|>


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

// Binary buttondemo shows the functionality of a button widget.
package main

import (
	"context"
	"fmt"
	"os/user"
	"time"

	"github.com/mum4k/termdash"
	"github.com/mum4k/termdash/align"
	"github.com/mum4k/termdash/cell"
	"github.com/mum4k/termdash/container"
	"github.com/mum4k/termdash/keyboard"
	"github.com/mum4k/termdash/linestyle"
	"github.com/mum4k/termdash/terminal/tcell"
	"github.com/mum4k/termdash/widgets/button"
	"github.com/mum4k/termdash/widgets/textinput"
)

// buttonChunks creates the text chunks for a button from the provided text.
func buttonChunks(text string) []*button.TextChunk {
	if len(text) == 0 {
		return nil
	}
	first := string(text[0])
	rest := string(text[1:])

	return []*button.TextChunk{
		button.NewChunk(
			"<",
			button.TextCellOpts(cell.FgColor(cell.ColorWhite)),
			button.FocusedTextCellOpts(cell.FgColor(cell.ColorBlack)),
			button.PressedTextCellOpts(cell.FgColor(cell.ColorBlack)),
		),
		button.NewChunk(
			first,
			button.TextCellOpts(cell.FgColor(cell.ColorRed)),
		),
		button.NewChunk(
			rest,
			button.TextCellOpts(cell.FgColor(cell.ColorWhite)),
			button.FocusedTextCellOpts(cell.FgColor(cell.ColorBlack)),
			button.PressedTextCellOpts(cell.FgColor(cell.ColorBlack)),
		),
		button.NewChunk(
			">",
			button.TextCellOpts(cell.FgColor(cell.ColorWhite)),
			button.FocusedTextCellOpts(cell.FgColor(cell.ColorBlack)),
			button.PressedTextCellOpts(cell.FgColor(cell.ColorBlack)),
		),
	}
}

func main() {
	t, err := tcell.New()
	if err != nil {
		panic(err)
	}
	defer t.Close()

	ctx, cancel := context.WithCancel(context.Background())

	var username string
	u, err := user.Current()
	if err != nil {
		username = "mum4k"
	} else {
		username = u.Username
	}

	userInput, err := textinput.New(
		textinput.Label("Username: ", cell.FgColor(cell.ColorNumber(33))),
		textinput.DefaultText(username),
		textinput.MaxWidthCells(20),
		textinput.ExclusiveKeyboardOnFocus(),
	)
	uidInput, err := textinput.New(
		textinput.Label("UID:      ", cell.FgColor(cell.ColorNumber(33))),
		textinput.DefaultText("1000"),
		textinput.MaxWidthCells(20),
		textinput.ExclusiveKeyboardOnFocus(),
	)
	gidInput, err := textinput.New(
		textinput.Label("GID:      ", cell.FgColor(cell.ColorNumber(33))),
		textinput.DefaultText("1000"),
		textinput.MaxWidthCells(20),
		textinput.ExclusiveKeyboardOnFocus(),
	)
	homeInput, err := textinput.New(
		textinput.Label("Home:     ", cell.FgColor(cell.ColorNumber(33))),
		textinput.DefaultText(fmt.Sprintf("/home/%s", username)),
		textinput.MaxWidthCells(20),
		textinput.ExclusiveKeyboardOnFocus(),
	)

	submitB, err := button.NewFromChunks(buttonChunks("Submit"), func() error {
		return nil
	},
		button.Key(keyboard.KeyEnter),
		button.GlobalKeys('s', 'S'),
		button.DisableShadow(),
		button.Height(1),
		button.TextHorizontalPadding(0),
		button.FillColor(cell.ColorBlack),
		button.FocusedFillColor(cell.ColorNumber(117)),
		button.PressedFillColor(cell.ColorNumber(220)),
	)
	if err != nil {
		panic(err)
	}

	cancelB, err := button.NewFromChunks(buttonChunks("Cancel"), func() error {
		cancel()
		return nil
	},
		button.FillColor(cell.ColorNumber(220)),
		button.Key(keyboard.KeyEnter),
		button.GlobalKeys('c', 'C'),
		button.DisableShadow(),
		button.Height(1),
		button.TextHorizontalPadding(0),
		button.FillColor(cell.ColorBlack),
		button.FocusedFillColor(cell.ColorNumber(117)),
		button.PressedFillColor(cell.ColorNumber(220)),
	)
	if err != nil {
		panic(err)
	}

	c, err := container.New(
		t,
		container.KeyFocusNext(keyboard.KeyTab),
		container.KeyFocusGroupsNext(keyboard.KeyArrowDown, 1),
		container.KeyFocusGroupsPrevious(keyboard.KeyArrowUp, 1),
		container.KeyFocusGroupsNext(keyboard.KeyArrowRight, 2),
		container.KeyFocusGroupsPrevious(keyboard.KeyArrowLeft, 2),
		container.SplitHorizontal(
			container.Top(
				container.Border(linestyle.Light),
				container.SplitHorizontal(
					container.Top(
						container.SplitHorizontal(
							container.Top(
								container.Focused(),
								container.KeyFocusGroups(1),
								container.PlaceWidget(userInput),
							),
							container.Bottom(
								container.KeyFocusGroups(1),
								container.KeyFocusSkip(),
								container.PlaceWidget(uidInput),
							),
						),
					),
					container.Bottom(
						container.SplitHorizontal(
							container.Top(
								container.KeyFocusGroups(1),
								container.KeyFocusSkip(),
								container.PlaceWidget(gidInput),
							),
							container.Bottom(
								container.KeyFocusGroups(1),
								container.KeyFocusSkip(),
								container.PlaceWidget(homeInput),
							),
						),
					),
				),
			),
			container.Bottom(
				container.SplitHorizontal(
					container.Top(
						container.SplitVertical(
							container.Left(
								container.KeyFocusGroups(1, 2),
								container.PlaceWidget(submitB),
								container.AlignHorizontal(align.HorizontalRight),
								container.PaddingRight(5),
							),
							container.Right(
								container.KeyFocusGroups(1, 2),
								container.PlaceWidget(cancelB),
								container.AlignHorizontal(align.HorizontalLeft),
								container.PaddingLeft(5),
							),
						),
					),
					container.Bottom(
						container.KeyFocusSkip(),
					),
					container.SplitFixed(3),
				),
			),
			container.SplitFixed(6),
		),
	)
	if err != nil {
		panic(err)
	}

	if err := termdash.Run(ctx, t, c, termdash.RedrawInterval(100*time.Millisecond)); err != nil {
		panic(err)
	}
}
<commit_msg>Refactor form construction.<commit_after>// Copyright 2019 Google Inc.
//
// 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.

// Binary formdemo creates a form that accepts text inputs and supports
// keyboard navigation.
package main

import (
	"context"
	"fmt"
	"os/user"
	"time"

	"github.com/mum4k/termdash"
	"github.com/mum4k/termdash/align"
	"github.com/mum4k/termdash/cell"
	"github.com/mum4k/termdash/container"
	"github.com/mum4k/termdash/keyboard"
	"github.com/mum4k/termdash/linestyle"
	"github.com/mum4k/termdash/terminal/tcell"
	"github.com/mum4k/termdash/widgets/button"
	"github.com/mum4k/termdash/widgets/textinput"
)

// buttonChunks creates the text chunks for a button from the provided text.
func buttonChunks(text string) []*button.TextChunk {
	if len(text) == 0 {
		return nil
	}
	first := string(text[0])
	rest := string(text[1:])

	return []*button.TextChunk{
		button.NewChunk(
			"<",
			button.TextCellOpts(cell.FgColor(cell.ColorWhite)),
			button.FocusedTextCellOpts(cell.FgColor(cell.ColorBlack)),
			button.PressedTextCellOpts(cell.FgColor(cell.ColorBlack)),
		),
		button.NewChunk(
			first,
			button.TextCellOpts(cell.FgColor(cell.ColorRed)),
		),
		button.NewChunk(
			rest,
			button.TextCellOpts(cell.FgColor(cell.ColorWhite)),
			button.FocusedTextCellOpts(cell.FgColor(cell.ColorBlack)),
			button.PressedTextCellOpts(cell.FgColor(cell.ColorBlack)),
		),
		button.NewChunk(
			">",
			button.TextCellOpts(cell.FgColor(cell.ColorWhite)),
			button.FocusedTextCellOpts(cell.FgColor(cell.ColorBlack)),
			button.PressedTextCellOpts(cell.FgColor(cell.ColorBlack)),
		),
	}
}

// form contains the elements of a text input form
type form struct {
	userInput *textinput.TextInput
	uidInput  *textinput.TextInput
	gidInput  *textinput.TextInput
	homeInput *textinput.TextInput
	submitB   *button.Button
	cancelB   *button.Button
}

// newForm returns a new form instance.
// The cancel argument is a function that terminates Termdash when called.
func newForm(cancel context.CancelFunc) (*form, error) {
	var username string
	u, err := user.Current()
	if err != nil {
		username = "mum4k"
	} else {
		username = u.Username
	}

	userInput, err := textinput.New(
		textinput.Label("Username: ", cell.FgColor(cell.ColorNumber(33))),
		textinput.DefaultText(username),
		textinput.MaxWidthCells(20),
		textinput.ExclusiveKeyboardOnFocus(),
	)
	uidInput, err := textinput.New(
		textinput.Label("UID:      ", cell.FgColor(cell.ColorNumber(33))),
		textinput.DefaultText("1000"),
		textinput.MaxWidthCells(20),
		textinput.ExclusiveKeyboardOnFocus(),
	)
	gidInput, err := textinput.New(
		textinput.Label("GID:      ", cell.FgColor(cell.ColorNumber(33))),
		textinput.DefaultText("1000"),
		textinput.MaxWidthCells(20),
		textinput.ExclusiveKeyboardOnFocus(),
	)
	homeInput, err := textinput.New(
		textinput.Label("Home:     ", cell.FgColor(cell.ColorNumber(33))),
		textinput.DefaultText(fmt.Sprintf("/home/%s", username)),
		textinput.MaxWidthCells(20),
		textinput.ExclusiveKeyboardOnFocus(),
	)

	submitB, err := button.NewFromChunks(buttonChunks("Submit"), func() error {
		return nil
	},
		button.Key(keyboard.KeyEnter),
		button.GlobalKeys('s', 'S'),
		button.DisableShadow(),
		button.Height(1),
		button.TextHorizontalPadding(0),
		button.FillColor(cell.ColorBlack),
		button.FocusedFillColor(cell.ColorNumber(117)),
		button.PressedFillColor(cell.ColorNumber(220)),
	)
	if err != nil {
		panic(err)
	}

	cancelB, err := button.NewFromChunks(buttonChunks("Cancel"), func() error {
		cancel()
		return nil
	},
		button.FillColor(cell.ColorNumber(220)),
		button.Key(keyboard.KeyEnter),
		button.GlobalKeys('c', 'C'),
		button.DisableShadow(),
		button.Height(1),
		button.TextHorizontalPadding(0),
		button.FillColor(cell.ColorBlack),
		button.FocusedFillColor(cell.ColorNumber(117)),
		button.PressedFillColor(cell.ColorNumber(220)),
	)
	if err != nil {
		panic(err)
	}

	return &form{
		userInput: userInput,
		uidInput:  uidInput,
		gidInput:  gidInput,
		homeInput: homeInput,
		submitB:   submitB,
		cancelB:   cancelB,
	}, nil
}

// formLayout updates the container into a layout with text inputs and buttons.
func formLayout(c *container.Container, f *form) error {
	return c.Update("root",
		container.KeyFocusNext(keyboard.KeyTab),
		container.KeyFocusGroupsNext(keyboard.KeyArrowDown, 1),
		container.KeyFocusGroupsPrevious(keyboard.KeyArrowUp, 1),
		container.KeyFocusGroupsNext(keyboard.KeyArrowRight, 2),
		container.KeyFocusGroupsPrevious(keyboard.KeyArrowLeft, 2),
		container.SplitHorizontal(
			container.Top(
				container.Border(linestyle.Light),
				container.SplitHorizontal(
					container.Top(
						container.SplitHorizontal(
							container.Top(
								container.Focused(),
								container.KeyFocusGroups(1),
								container.PlaceWidget(f.userInput),
							),
							container.Bottom(
								container.KeyFocusGroups(1),
								container.KeyFocusSkip(),
								container.PlaceWidget(f.uidInput),
							),
						),
					),
					container.Bottom(
						container.SplitHorizontal(
							container.Top(
								container.KeyFocusGroups(1),
								container.KeyFocusSkip(),
								container.PlaceWidget(f.gidInput),
							),
							container.Bottom(
								container.KeyFocusGroups(1),
								container.KeyFocusSkip(),
								container.PlaceWidget(f.homeInput),
							),
						),
					),
				),
			),
			container.Bottom(
				container.SplitHorizontal(
					container.Top(
						container.SplitVertical(
							container.Left(
								container.KeyFocusGroups(1, 2),
								container.PlaceWidget(f.submitB),
								container.AlignHorizontal(align.HorizontalRight),
								container.PaddingRight(5),
							),
							container.Right(
								container.KeyFocusGroups(1, 2),
								container.PlaceWidget(f.cancelB),
								container.AlignHorizontal(align.HorizontalLeft),
								container.PaddingLeft(5),
							),
						),
					),
					container.Bottom(
						container.KeyFocusSkip(),
					),
					container.SplitFixed(3),
				),
			),
			container.SplitFixed(6),
		),
	)
}

func main() {
	t, err := tcell.New()
	if err != nil {
		panic(err)
	}
	defer t.Close()

	ctx, cancel := context.WithCancel(context.Background())
	c, err := container.New(t, container.ID("root"))
	if err != nil {
		panic(err)
	}

	f, err := newForm(cancel)
	if err != nil {
		panic(err)
	}
	if err := formLayout(c, f); err != nil {
		panic(err)
	}

	if err := termdash.Run(ctx, t, c, termdash.RedrawInterval(100*time.Millisecond)); err != nil {
		panic(err)
	}
}
<|endoftext|>


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

// +build linux

package diskmanager

import (
	"bufio"
	"bytes"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"syscall"

	"github.com/juju/errors"

	"github.com/juju/juju/storage"
)

var pairsRE = regexp.MustCompile(`([A-Z]+)=(?:"(.*?)")`)

const (
	// values for the TYPE column that we care about

	typeDisk = "disk"
	typeLoop = "loop"
)

func init() {
	DefaultListBlockDevices = listBlockDevices
}

func listBlockDevices() ([]storage.BlockDevice, error) {
	columns := []string{
		"KNAME",      // kernel name
		"SIZE",       // size
		"LABEL",      // filesystem label
		"UUID",       // filesystem UUID
		"FSTYPE",     // filesystem type
		"TYPE",       // device type
		"MOUNTPOINT", // moint point
	}

	logger.Debugf("executing lsblk")
	output, err := exec.Command(
		"lsblk",
		"-b", // output size in bytes
		"-P", // output fields as key=value pairs
		"-o", strings.Join(columns, ","),
	).Output()
	if err != nil {
		return nil, errors.Annotate(
			err, "cannot list block devices: lsblk failed",
		)
	}

	blockDeviceMap := make(map[string]storage.BlockDevice)
	s := bufio.NewScanner(bytes.NewReader(output))
	for s.Scan() {
		pairs := pairsRE.FindAllStringSubmatch(s.Text(), -1)
		var dev storage.BlockDevice
		var deviceType string
		for _, pair := range pairs {
			switch pair[1] {
			case "KNAME":
				dev.DeviceName = pair[2]
			case "SIZE":
				size, err := strconv.ParseUint(pair[2], 10, 64)
				if err != nil {
					logger.Errorf(
						"invalid size %q from lsblk: %v", pair[2], err,
					)
				} else {
					dev.Size = size / bytesInMiB
				}
			case "LABEL":
				dev.Label = pair[2]
			case "UUID":
				dev.UUID = pair[2]
			case "FSTYPE":
				dev.FilesystemType = pair[2]
			case "TYPE":
				deviceType = pair[2]
			case "MOUNTPOINT":
				dev.MountPoint = pair[2]
			default:
				logger.Debugf("unexpected field from lsblk: %q", pair[1])
			}
		}

		// We may later want to expand this, e.g. to handle lvm,
		// dmraid, crypt, etc., but this is enough to cover bases
		// for now.
		switch deviceType {
		case typeDisk, typeLoop:
		default:
			logger.Debugf("ignoring %q type device: %+v", deviceType, dev)
			continue
		}

		// Check if the block device is in use. We need to know this so we can
		// issue an error if the user attempts to allocate an in-use disk to a
		// unit.
		dev.InUse, err = blockDeviceInUse(dev)
		if os.IsNotExist(err) {
			// In LXC containers, lsblk will show the block devices of the
			// host, but the devices will typically not be present.
			continue
		} else if err != nil {
			logger.Errorf(
				"error checking if %q is in use: %v", dev.DeviceName, err,
			)
			// We cannot detect, so err on the side of caution and default to
			// "in use" so the device cannot be used.
			dev.InUse = true
		}
		blockDeviceMap[dev.DeviceName] = dev
	}
	if err := s.Err(); err != nil {
		return nil, errors.Annotate(err, "cannot parse lsblk output")
	}

	blockDevices := make([]storage.BlockDevice, 0, len(blockDeviceMap))
	for _, dev := range blockDeviceMap {
		blockDevices = append(blockDevices, dev)
	}
	return blockDevices, nil
}

// blockDeviceInUse checks if the specified block device
// is in use by attempting to open the device exclusively.
//
// If the error returned satisfies os.IsNotExists, then
// the device will be ignored altogether.
var blockDeviceInUse = func(dev storage.BlockDevice) (bool, error) {
	f, err := os.OpenFile("/dev/"+dev.DeviceName, os.O_EXCL, 0)
	if err == nil {
		f.Close()
		return false, nil
	}
	perr, ok := err.(*os.PathError)
	if !ok {
		return false, err
	}
	// open(2): "In general, the behavior of O_EXCL is undefined if
	// it is used without O_CREAT. There is one exception: on Linux
	// 2.6 and later, O_EXCL can be used without O_CREAT if pathname
	// refers to a block device. If the block device is in use by the
	// system  (e.g., mounted), open() fails with the error EBUSY."
	if errno, _ := perr.Err.(syscall.Errno); errno == syscall.EBUSY {
		return true, nil
	}
	return false, err
}
<commit_msg>worker/diskmanager: Tracef ignored devices<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// +build linux

package diskmanager

import (
	"bufio"
	"bytes"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"syscall"

	"github.com/juju/errors"

	"github.com/juju/juju/storage"
)

var pairsRE = regexp.MustCompile(`([A-Z]+)=(?:"(.*?)")`)

const (
	// values for the TYPE column that we care about

	typeDisk = "disk"
	typeLoop = "loop"
)

func init() {
	DefaultListBlockDevices = listBlockDevices
}

func listBlockDevices() ([]storage.BlockDevice, error) {
	columns := []string{
		"KNAME",      // kernel name
		"SIZE",       // size
		"LABEL",      // filesystem label
		"UUID",       // filesystem UUID
		"FSTYPE",     // filesystem type
		"TYPE",       // device type
		"MOUNTPOINT", // moint point
	}

	logger.Debugf("executing lsblk")
	output, err := exec.Command(
		"lsblk",
		"-b", // output size in bytes
		"-P", // output fields as key=value pairs
		"-o", strings.Join(columns, ","),
	).Output()
	if err != nil {
		return nil, errors.Annotate(
			err, "cannot list block devices: lsblk failed",
		)
	}

	blockDeviceMap := make(map[string]storage.BlockDevice)
	s := bufio.NewScanner(bytes.NewReader(output))
	for s.Scan() {
		pairs := pairsRE.FindAllStringSubmatch(s.Text(), -1)
		var dev storage.BlockDevice
		var deviceType string
		for _, pair := range pairs {
			switch pair[1] {
			case "KNAME":
				dev.DeviceName = pair[2]
			case "SIZE":
				size, err := strconv.ParseUint(pair[2], 10, 64)
				if err != nil {
					logger.Errorf(
						"invalid size %q from lsblk: %v", pair[2], err,
					)
				} else {
					dev.Size = size / bytesInMiB
				}
			case "LABEL":
				dev.Label = pair[2]
			case "UUID":
				dev.UUID = pair[2]
			case "FSTYPE":
				dev.FilesystemType = pair[2]
			case "TYPE":
				deviceType = pair[2]
			case "MOUNTPOINT":
				dev.MountPoint = pair[2]
			default:
				logger.Debugf("unexpected field from lsblk: %q", pair[1])
			}
		}

		// We may later want to expand this, e.g. to handle lvm,
		// dmraid, crypt, etc., but this is enough to cover bases
		// for now.
		switch deviceType {
		case typeDisk, typeLoop:
		default:
			logger.Tracef("ignoring %q type device: %+v", deviceType, dev)
			continue
		}

		// Check if the block device is in use. We need to know this so we can
		// issue an error if the user attempts to allocate an in-use disk to a
		// unit.
		dev.InUse, err = blockDeviceInUse(dev)
		if os.IsNotExist(err) {
			// In LXC containers, lsblk will show the block devices of the
			// host, but the devices will typically not be present.
			continue
		} else if err != nil {
			logger.Errorf(
				"error checking if %q is in use: %v", dev.DeviceName, err,
			)
			// We cannot detect, so err on the side of caution and default to
			// "in use" so the device cannot be used.
			dev.InUse = true
		}
		blockDeviceMap[dev.DeviceName] = dev
	}
	if err := s.Err(); err != nil {
		return nil, errors.Annotate(err, "cannot parse lsblk output")
	}

	blockDevices := make([]storage.BlockDevice, 0, len(blockDeviceMap))
	for _, dev := range blockDeviceMap {
		blockDevices = append(blockDevices, dev)
	}
	return blockDevices, nil
}

// blockDeviceInUse checks if the specified block device
// is in use by attempting to open the device exclusively.
//
// If the error returned satisfies os.IsNotExists, then
// the device will be ignored altogether.
var blockDeviceInUse = func(dev storage.BlockDevice) (bool, error) {
	f, err := os.OpenFile("/dev/"+dev.DeviceName, os.O_EXCL, 0)
	if err == nil {
		f.Close()
		return false, nil
	}
	perr, ok := err.(*os.PathError)
	if !ok {
		return false, err
	}
	// open(2): "In general, the behavior of O_EXCL is undefined if
	// it is used without O_CREAT. There is one exception: on Linux
	// 2.6 and later, O_EXCL can be used without O_CREAT if pathname
	// refers to a block device. If the block device is in use by the
	// system  (e.g., mounted), open() fails with the error EBUSY."
	if errno, _ := perr.Err.(syscall.Errno); errno == syscall.EBUSY {
		return true, nil
	}
	return false, err
}
<|endoftext|>


<commit_before>package dft

import (
	"github.com/huangjunwen/JustSQL/render"
)

func init() {
	render.RegistBuiltinTemplate("table", render.DefaultTemplateSetName, `
{{/* =========================== */}}
{{/*          imports            */}}
{{/* =========================== */}}
{{- $ctx := imp "context" -}}
{{- $fmt := imp "fmt" -}}
{{- $sql := imp "database/sql" -}}
{{- $driver := imp "database/sql/driver" -}}
{{- $sqlx := imp "github.com/jmoiron/sqlx" -}}
{{- $strings := imp "strings" -}}

{{/* =========================== */}}
{{/*      global variables       */}}
{{/* =========================== */}}
{{- $tableName := .Table.Name -}}
{{- $structName := .Table.PascalName -}}
{{- $cols := .Table.Columns -}}
{{- $autoIncCol := .Table.AutoIncColumn -}}
{{- $primaryCols := .Table.PrimaryColumns -}}

{{- $structFieldNameList := stringList -}}
{{- $structFieldTypeList := stringList -}}
{{- $enumColList := columnList -}}
{{- $setColList := columnList -}}
{{- range $i, $col := $cols -}}
	{{- append $structFieldNameList $col.PascalName -}}
	{{- if $col.IsEnum -}}
		{{- append $structFieldTypeList (printf "%s%s" $structName $col.PascalName) -}}
		{{- append $enumColList $col -}}
	{{- else if $col.IsSet -}}
		{{- append $structFieldTypeList (printf "%s%s" $structName $col.PascalName) -}}
		{{- append $setColList $col -}}
	{{- else -}}
		{{- append $structFieldTypeList (typeName $col.Type) -}}
	{{- end -}}
{{- end -}}
{{- $structFieldNames := $structFieldNameList.Strings -}}
{{- $structFieldTypes := $structFieldTypeList.Strings -}}
{{- $enumCols := $enumColList.Cols -}}
{{- $setCols := $setColList.Cols -}}

{{/* =========================== */}}
{{/*          enums              */}}
{{/* =========================== */}}

{{ range $i, $col := $enumCols }}
	{{/* =========================== */}}
	{{/*        enum variables       */}}
	{{/* =========================== */}}
	{{- $enumName := printf "%s%s" $structName $col.PascalName -}}
	{{- $enumItemList := stringList -}}
	{{- range $i, $item := $col.Elems -}}
		{{- if eq $item "" -}}
			{{- append $enumItemList (printf "%sEmpty_" $enumName) -}}
		{{- else -}}
			{{- append $enumItemList (printf "%s%s" $enumName (pascal $item)) }}
		{{- end -}}
	{{- end -}}
	{{- $enumItems := $enumItemList.Strings -}}

	{{/* =========================== */}}
	{{/*        enum code            */}}
	{{/* =========================== */}}

// Enum {{ $enumName }}.
type {{ $enumName }} int

// Enum {{ $enumName }} items.
const (
	// NULL value.
	{{ $enumName }}NULL_ = {{ $enumName }}(0)
	{{- range $i, $item := $enumItems }}
	// {{ printf "%+q" (index $col.Elems $i) }}
	{{ $item }} = {{ $enumName }}({{ $i }})
	{{- end }}
)

func (e {{ $enumName }}) String() string {
	switch e {
	{{- range $i, $item := $enumItems }}
	case {{ $item }}:
		return {{ printf "%+q" (index $col.Elems $i) }}
	{{- end }}
	}
	return ""
}

func (e {{ $enumName }}) Valid() bool {
	return int(e) > 0 && int(e) <= {{ printf "%d" (len $col.Elems) }}
}

// Scan implements database/sql.Scanner interface.
func (e *{{ $enumName }}) Scan(value interface{}) error {
	if value == nil {
		*e = {{ $enumName }}NULL_
		return nil
	}

	switch s := string(value.([]byte)); s {
	{{- range $i, $item := $enumItems }}
	case {{ printf "%+q" (index $col.Elems $i) }}:
		*e = {{ $item }}
	{{- end }}
	default:
		return {{ $fmt }}.Errorf("Unexpected value for {{ $enumName }}: %+q", s)
	}

	return nil
}

// Value implements database/sql/driver.Valuer interface.
func (e {{ $enumName }}) Value() (driver.Value, error) {
	if !e.Valid() {
		return nil, nil
	}
	return e.String(), nil
}

{{ end }}

{{/* =========================== */}}
{{/*          sets               */}}
{{/* =========================== */}}

{{ range $i, $col := $setCols -}}

	{{/* =========================== */}}
	{{/*        set  variables       */}}
	{{/* =========================== */}}
	{{- $setName := printf "%s%s" $structName $col.PascalName -}}
	{{- $setItemList := stringList -}}
	{{- range $i, $item := $col.Elems -}}
		{{- if eq $item "" -}}
			{{- append $setItemList (printf "%sEmpty_" $setName) -}}
		{{- else -}}
			{{- append $setItemList (printf "%s%s" $setName (pascal $item)) }}
		{{- end -}}
	{{- end -}}
	{{- $setItems := $setItemList.Strings -}}

	{{/* =========================== */}}
	{{/*        set  code            */}}
	{{/* =========================== */}}

// Set {{ $setName }}.
type {{ $setName }} struct {
	val uint64 // Up to 64 distinct members. See https://dev.mysql.com/doc/refman/5.7/en/set.html
	valid bool // NULL if valid is false.
}

// Set {{ $setName }} items.
const (
	{{- range $i, $item := $setItems }}
	// {{ printf "%+q" (index $col.Elems $i) }}
	{{ $item }} = uint64(1<<{{ $i }})
	{{- end }}
)

func New{{ $setName }}(items ...uint64) {{ $setName }} {
	var val uint64 = 0
	for _, item := range items {
		if item > 0 && (item & (item - 1)) == 0 && item <= (1 << ({{ len $col.Elems }} - 1)) {
			val |= item
		}
	}
	return {{ $setName }}{
		val: val,
		valid: true,
	}
}

func (s {{ $setName }}) String() string {
	parts := make([]string, 0)
	{{- range $i, $item:= $setItems }}
	if s.val & {{ $item }} != 0 {
		parts = append(parts, {{ printf "%+q" (index $col.Elems $i) }})
	}
	{{- end }}
	return strings.Join(parts, ",")
}

func (s {{ $setName }}) Valid() bool {
	return s.valid
}

// Scan implements database/sql.Scanner interface.
func (s *{{ $setName }}) Scan(value interface{}) error {
	if value == nil {
		s.val = 0
		s.valid = false
		return nil
	}

	var val uint64 = 0
	for _, part := range {{ $strings }}.Split(string(value.([]byte)), ",") {
		switch part {
	{{- range $i, $elem := $col.Elems }}
		case {{ printf "%+q" $elem }}:
			val |= {{ index $setItems $i }}
	{{- end }}
		default:
			return {{ $fmt }}.Errorf("Unexpected value for {{ $setName }}: %+q", part)
		}
	}

	s.val = val
	s.valid = true
	return nil
}

// Value implements database/sql/driver.Valuer interface.
func (s {{ $setName }}) Value() (driver.Value, error) {
	if !s.Valid() {
		return nil, nil
	}
	return s.String(), nil
}

{{- end }}

{{/* =========================== */}}
{{/*          main struct        */}}
{{/* =========================== */}}

// {{ $structName }} represents an entry of table "{{ $tableName }}".
type {{ $structName }} struct {
{{- range $i, $col := $cols }}
	{{ index $structFieldNames $i }} {{ index $structFieldTypes $i }} `+"`db:\"{{ $col.Name }}\"`"+` // {{ $col.Name }}
{{- end }}
}

{{/* =========================== */}}
{{/*     insert/update/delete    */}}
{{/* =========================== */}}

// Insert insert an entry of {{ $tableName }} into database.
func (entry_ *{{ $structName }}) Insert(ctx_ {{ $ctx }}.Context, db_ DBer) error {

	sql_ := {{ $sqlx }}.Rebind(BindType, "INSERT INTO {{ $tableName }} " +
		"({{ join (columnNames $cols) ", " }}) " +
		"VALUES ({{ join (dup "?" (len $cols))  ", " }})")

	{{ if notNil $autoIncCol }}res_{{ else }}_{{ end }}, err_ := db_.ExecContext(ctx_, sql_{{ range $i, $field := $structFieldNames }}, entry_.{{ $field }}{{ end }})
	if err_ != nil {
		return err_
	}

	{{ if notNil $autoIncCol -}}
	lastInsertId_, err_ := res_.LastInsertId()
	if err_ != nil {
		return err_
	}

	return SaveCoerceFromInt64(lastInsertId_, &entry_.{{ $autoIncCol.PascalName }})
	{{ else -}}
	return nil
	{{ end -}}
}

{{ if ne (len $primaryCols) 0 -}}

func (entry_ *{{ $structName }}) Update(ctx_ {{ $ctx }}.Context, db_ DBer) (int64, error) {

	sql_ := {{ $sqlx }}.Rebind(BindType, "UPDATE {{ $tableName }} " +
		"SET {{ range $i, $col := $cols }}{{ if ne $i 0 }}, {{ end }}{{ $col.Name }}=?{{ end }} " +
		"WHERE {{ range $i, $col := $primaryCols }}{{ if ne $i 0 }}AND {{ end }}{{ $col.Name }}=? {{ end }}")

	r_, err_ := db_.ExecContext(ctx_, sql_{{ range $i, $col := $cols }}, entry_.{{ $col.PascalName }}{{ end }}{{ range $i, $col := $primaryCols }}, entry_.{{ $col.PascalName }}{{ end }})
	if err_ != nil {
		return 0, err_
	}

	return r_.RowsAffected()

}

func (entry_ *{{ $structName }}) Delete(ctx_ {{ $ctx }}.Context, db_ DBer) (int64, error) {

	sql_ := {{ $sqlx }}.Rebind(BindType, "DELETE FROM {{ $tableName }} " +
		"WHERE {{ range $i, $col := $primaryCols }}{{ if ne $i 0 }}AND {{ end }}{{ $col.Name }}=? {{ end }}")

	r_, err_ := db_.ExecContext(ctx_, sql_{{ range $i, $col := $primaryCols }}, entry_.{{ $col.PascalName }}{{ end }})
	if err_ != nil {
		return 0, err_
	}

	return r_.RowsAffected()
}

{{ end }}

{{/* =========================== */}}
{{/*         foreign key         */}}
{{/* =========================== */}}

{{- range $i, $fk := .Table.ForeignKeys }}
	{{- $refIndex := $fk.RefIndex }}
	{{- if $refIndex.Unique }}
	{{/* =========================== */}}
	{{/*    foreign key variables    */}}
	{{/* =========================== */}}
		{{- $fkColumns := $fk.Columns }}
		{{- $refTable := $fk.RefTable }}
		{{- $refColumns := $fk.RefColumns }}

	{{/* =========================== */}}
	{{/*       foreign key code      */}}
	{{/* =========================== */}}

// {{ $refTable.PascalName }} return {{ printf "%q" $refTable.Name }} entry by foreign key "{{ printf "%s.%s" $fk.Table.Name $fk.Name }}".
func (entry_ *{{ $structName }}) {{ $refTable.PascalName }}(ctx_ {{ $ctx }}.Context, db_ DBer) (*{{ $refTable.PascalName }}, error) {
	return {{ $refTable.PascalName }}By{{ $refIndex.PascalName }}(ctx_, db_
		{{- range $j, $fkCol := $fkColumns -}}
			{{- $expr := printf "entry_.%s" $fkCol.PascalName -}}
			{{- $fkColType := typeName $fkCol -}}
			{{- $refColType := typeName (index $refColumns $j) -}}
		, {{ cast $expr $fkColType $refColType }}
		{{- end -}})
}
	{{- end }}
{{- end }}

{{/* =========================== */}}
{{/*          unique indices     */}}
{{/* =========================== */}}

{{- range $i, $index := .Table.Indices }}
	{{- if $index.Unique }}
	{{/* =========================== */}}
	{{/*   unique index variables    */}}
	{{/* =========================== */}}
		{{- $indexCols := $index.Columns }}
		{{- $argNameList := stringList }}
		{{- $argTypeList := stringList }}
		{{- range $j, $col := $indexCols }}
			{{- append $argNameList (camel $col.Name) }}
			{{- append $argTypeList (typeName $col.Type) }}
		{{- end }}
		{{- $argNames := $argNameList.Strings }}
		{{- $argTypes := $argTypeList.Strings }}

	{{/* =========================== */}}
	{{/*      unique index code      */}}
	{{/* =========================== */}}

// {{ $structName }}By{{ $index.PascalName }} query {{ printf "%+q" $tableName }} table by {{ if $index.Primary }}primary key{{ else }}unique key {{ printf "%+q" $index.Name }}{{ end }}.
// Return nil if error occurred or there is not row found.
func {{ $structName }}By{{ $index.PascalName }}(ctx_ {{ $ctx }}.Context, db_ DBer{{ range $j, $argName := $argNames }}, {{ $argName }} {{ index $argTypes $j }}{{ end }}) (*{{ $structName }}, error) {

	sql_ := {{ $sqlx }}.Rebind(BindType, "SELECT {{ join (columnNames $cols) ", " }} " +
		"FROM {{ $tableName }} " +
		"WHERE {{ range $j, $col := $indexCols }}{{ if ne $j 0 }}AND {{ end }}{{ $col.Name }}=? {{ end }}")

	row_ := db_.QueryRowContext(ctx_, sql_{{ range $j, $argName := $argNames }}, {{ $argName }}{{ end }})

	entry_ := new({{ $structName }})
	if err_ := row_.Scan({{ range $j, $field := $structFieldNames }}{{ if ne $j 0 }}, {{ end }}&entry_.{{ $field }}{{ end }}); err_ != nil {
		if err_ == {{ $sql }}.ErrNoRows {
			return nil, nil
		}
		return nil, err_
	}

	return entry_, nil
}

	{{- end }}
{{- end }}

`)

}
<commit_msg>Fix table template bugs<commit_after>package dft

import (
	"github.com/huangjunwen/JustSQL/render"
)

func init() {
	render.RegistBuiltinTemplate("table", render.DefaultTemplateSetName, `
{{/* =========================== */}}
{{/*          imports            */}}
{{/* =========================== */}}
{{- $ctx := imp "context" -}}
{{- $fmt := imp "fmt" -}}
{{- $sql := imp "database/sql" -}}
{{- $driver := imp "database/sql/driver" -}}
{{- $sqlx := imp "github.com/jmoiron/sqlx" -}}
{{- $strings := imp "strings" -}}

{{/* =========================== */}}
{{/*      global variables       */}}
{{/* =========================== */}}
{{- $tableName := .Table.Name -}}
{{- $structName := .Table.PascalName -}}
{{- $cols := .Table.Columns -}}
{{- $autoIncCol := .Table.AutoIncColumn -}}
{{- $primaryCols := .Table.PrimaryColumns -}}

{{- $structFieldNameList := stringList -}}
{{- $structFieldTypeList := stringList -}}
{{- $enumColList := columnList -}}
{{- $setColList := columnList -}}
{{- range $i, $col := $cols -}}
	{{- append $structFieldNameList $col.PascalName -}}
	{{- if $col.IsEnum -}}
		{{- append $structFieldTypeList (printf "%s%s" $structName $col.PascalName) -}}
		{{- append $enumColList $col -}}
	{{- else if $col.IsSet -}}
		{{- append $structFieldTypeList (printf "%s%s" $structName $col.PascalName) -}}
		{{- append $setColList $col -}}
	{{- else -}}
		{{- append $structFieldTypeList (typeName $col.Type) -}}
	{{- end -}}
{{- end -}}
{{- $structFieldNames := $structFieldNameList.Strings -}}
{{- $structFieldTypes := $structFieldTypeList.Strings -}}
{{- $enumCols := $enumColList.Cols -}}
{{- $setCols := $setColList.Cols -}}

{{/* =========================== */}}
{{/*          enums              */}}
{{/* =========================== */}}

{{ range $i, $col := $enumCols }}
	{{/* =========================== */}}
	{{/*        enum variables       */}}
	{{/* =========================== */}}
	{{- $enumName := printf "%s%s" $structName $col.PascalName -}}
	{{- $enumItemList := stringList -}}
	{{- range $i, $item := $col.Elems -}}
		{{- if eq $item "" -}}
			{{- append $enumItemList (printf "%sEmpty_" $enumName) -}}
		{{- else -}}
			{{- append $enumItemList (printf "%s%s" $enumName (pascal $item)) }}
		{{- end -}}
	{{- end -}}
	{{- $enumItems := $enumItemList.Strings -}}

	{{/* =========================== */}}
	{{/*        enum code            */}}
	{{/* =========================== */}}

// Enum {{ $enumName }}.
type {{ $enumName }} int

// Enum {{ $enumName }} items.
const (
	// NULL value.
	{{ $enumName }}NULL_ = {{ $enumName }}(0)
	{{- range $i, $item := $enumItems }}
	// {{ printf "%+q" (index $col.Elems $i) }}
	{{ $item }} = {{ $enumName }}({{ $i }} + 1)
	{{- end }}
)

func (e {{ $enumName }}) String() string {
	switch e {
	{{- range $i, $item := $enumItems }}
	case {{ $item }}:
		return {{ printf "%+q" (index $col.Elems $i) }}
	{{- end }}
	}
	return ""
}

func (e {{ $enumName }}) Valid() bool {
	return int(e) > 0 && int(e) <= {{ printf "%d" (len $col.Elems) }}
}

// Scan implements database/sql.Scanner interface.
func (e *{{ $enumName }}) Scan(value interface{}) error {
	if value == nil {
		*e = {{ $enumName }}NULL_
		return nil
	}

	switch s := string(value.([]byte)); s {
	{{- range $i, $item := $enumItems }}
	case {{ printf "%+q" (index $col.Elems $i) }}:
		*e = {{ $item }}
	{{- end }}
	default:
		return {{ $fmt }}.Errorf("Unexpected value for {{ $enumName }}: %+q", s)
	}

	return nil
}

// Value implements database/sql/driver.Valuer interface.
func (e {{ $enumName }}) Value() ({{ $driver }}.Value, error) {
	if !e.Valid() {
		return nil, nil
	}
	return e.String(), nil
}

{{ end }}

{{/* =========================== */}}
{{/*          sets               */}}
{{/* =========================== */}}

{{ range $i, $col := $setCols -}}

	{{/* =========================== */}}
	{{/*        set  variables       */}}
	{{/* =========================== */}}
	{{- $setName := printf "%s%s" $structName $col.PascalName -}}
	{{- $setItemList := stringList -}}
	{{- range $i, $item := $col.Elems -}}
		{{- if eq $item "" -}}
			{{- append $setItemList (printf "%sEmpty_" $setName) -}}
		{{- else -}}
			{{- append $setItemList (printf "%s%s" $setName (pascal $item)) }}
		{{- end -}}
	{{- end -}}
	{{- $setItems := $setItemList.Strings -}}

	{{/* =========================== */}}
	{{/*        set  code            */}}
	{{/* =========================== */}}

// Set {{ $setName }}.
type {{ $setName }} struct {
	val uint64 // Up to 64 distinct members. See https://dev.mysql.com/doc/refman/5.7/en/set.html
	valid bool // NULL if valid is false.
}

// Set {{ $setName }} items.
const (
	{{- range $i, $item := $setItems }}
	// {{ printf "%+q" (index $col.Elems $i) }}
	{{ $item }} = uint64(1<<{{ $i }})
	{{- end }}
)

func New{{ $setName }}(items ...uint64) {{ $setName }} {
	var val uint64 = 0
	for _, item := range items {
		if item > 0 && (item & (item - 1)) == 0 && item <= (1 << ({{ len $col.Elems }} - 1)) {
			val |= item
		}
	}
	return {{ $setName }}{
		val: val,
		valid: true,
	}
}

func (s {{ $setName }}) String() string {
	parts := make([]string, 0)
	{{- range $i, $item:= $setItems }}
	if s.val & {{ $item }} != 0 {
		parts = append(parts, {{ printf "%+q" (index $col.Elems $i) }})
	}
	{{- end }}
	return strings.Join(parts, ",")
}

func (s {{ $setName }}) Valid() bool {
	return s.valid
}

// Scan implements database/sql.Scanner interface.
func (s *{{ $setName }}) Scan(value interface{}) error {
	if value == nil {
		s.val = 0
		s.valid = false
		return nil
	}

	var val uint64 = 0
	for _, part := range {{ $strings }}.Split(string(value.([]byte)), ",") {
		switch part {
	{{- range $i, $elem := $col.Elems }}
		case {{ printf "%+q" $elem }}:
			val |= {{ index $setItems $i }}
	{{- end }}
		default:
			return {{ $fmt }}.Errorf("Unexpected value for {{ $setName }}: %+q", part)
		}
	}

	s.val = val
	s.valid = true
	return nil
}

// Value implements database/sql/driver.Valuer interface.
func (s {{ $setName }}) Value() (driver.Value, error) {
	if !s.Valid() {
		return nil, nil
	}
	return s.String(), nil
}

{{- end }}

{{/* =========================== */}}
{{/*          main struct        */}}
{{/* =========================== */}}

// {{ $structName }} represents an entry of table "{{ $tableName }}".
type {{ $structName }} struct {
{{- range $i, $col := $cols }}
	{{ index $structFieldNames $i }} {{ index $structFieldTypes $i }} `+"`db:\"{{ $col.Name }}\"`"+` // {{ $col.Name }}
{{- end }}
}

{{/* =========================== */}}
{{/*     insert/update/delete    */}}
{{/* =========================== */}}

// Insert insert an entry of {{ $tableName }} into database.
func (entry_ *{{ $structName }}) Insert(ctx_ {{ $ctx }}.Context, db_ DBer) error {

	sql_ := {{ $sqlx }}.Rebind(BindType, "INSERT INTO {{ $tableName }} " +
		"({{ join (columnNames $cols) ", " }}) " +
		"VALUES ({{ join (dup "?" (len $cols))  ", " }})")

	{{ if notNil $autoIncCol }}res_{{ else }}_{{ end }}, err_ := db_.ExecContext(ctx_, sql_{{ range $i, $field := $structFieldNames }}, entry_.{{ $field }}{{ end }})
	if err_ != nil {
		return err_
	}

	{{ if notNil $autoIncCol -}}
	lastInsertId_, err_ := res_.LastInsertId()
	if err_ != nil {
		return err_
	}

	return SaveCoerceFromInt64(lastInsertId_, &entry_.{{ $autoIncCol.PascalName }})
	{{ else -}}
	return nil
	{{ end -}}
}

{{ if ne (len $primaryCols) 0 -}}

func (entry_ *{{ $structName }}) Update(ctx_ {{ $ctx }}.Context, db_ DBer) (int64, error) {

	sql_ := {{ $sqlx }}.Rebind(BindType, "UPDATE {{ $tableName }} " +
		"SET {{ range $i, $col := $cols }}{{ if ne $i 0 }}, {{ end }}{{ $col.Name }}=?{{ end }} " +
		"WHERE {{ range $i, $col := $primaryCols }}{{ if ne $i 0 }}AND {{ end }}{{ $col.Name }}=? {{ end }}")

	r_, err_ := db_.ExecContext(ctx_, sql_{{ range $i, $col := $cols }}, entry_.{{ $col.PascalName }}{{ end }}{{ range $i, $col := $primaryCols }}, entry_.{{ $col.PascalName }}{{ end }})
	if err_ != nil {
		return 0, err_
	}

	return r_.RowsAffected()

}

func (entry_ *{{ $structName }}) Delete(ctx_ {{ $ctx }}.Context, db_ DBer) (int64, error) {

	sql_ := {{ $sqlx }}.Rebind(BindType, "DELETE FROM {{ $tableName }} " +
		"WHERE {{ range $i, $col := $primaryCols }}{{ if ne $i 0 }}AND {{ end }}{{ $col.Name }}=? {{ end }}")

	r_, err_ := db_.ExecContext(ctx_, sql_{{ range $i, $col := $primaryCols }}, entry_.{{ $col.PascalName }}{{ end }})
	if err_ != nil {
		return 0, err_
	}

	return r_.RowsAffected()
}

{{ end }}

{{/* =========================== */}}
{{/*         foreign key         */}}
{{/* =========================== */}}

{{- range $i, $fk := .Table.ForeignKeys }}
	{{- $refIndex := $fk.RefIndex }}
	{{- if $refIndex.Unique }}
	{{/* =========================== */}}
	{{/*    foreign key variables    */}}
	{{/* =========================== */}}
		{{- $fkColumns := $fk.Columns }}
		{{- $refTable := $fk.RefTable }}
		{{- $refColumns := $fk.RefColumns }}

	{{/* =========================== */}}
	{{/*       foreign key code      */}}
	{{/* =========================== */}}

// {{ $refTable.PascalName }} return {{ printf "%q" $refTable.Name }} entry by foreign key "{{ printf "%s.%s" $fk.Table.Name $fk.Name }}".
func (entry_ *{{ $structName }}) {{ $refTable.PascalName }}(ctx_ {{ $ctx }}.Context, db_ DBer) (*{{ $refTable.PascalName }}, error) {
	return {{ $refTable.PascalName }}By{{ $refIndex.PascalName }}(ctx_, db_
		{{- range $j, $fkCol := $fkColumns -}}
			{{- $expr := printf "entry_.%s" $fkCol.PascalName -}}
			{{- $fkColType := typeName $fkCol -}}
			{{- $refColType := typeName (index $refColumns $j) -}}
		, {{ cast $expr $fkColType $refColType }}
		{{- end -}})
}
	{{- end }}
{{- end }}

{{/* =========================== */}}
{{/*          unique indices     */}}
{{/* =========================== */}}

{{- range $i, $index := .Table.Indices }}
	{{- if $index.Unique }}
	{{/* =========================== */}}
	{{/*   unique index variables    */}}
	{{/* =========================== */}}
		{{- $indexCols := $index.Columns }}
		{{- $argNameList := stringList }}
		{{- $argTypeList := stringList }}
		{{- range $j, $col := $indexCols }}
			{{- append $argNameList (camel $col.Name) }}
			{{- append $argTypeList (typeName $col.Type) }}
		{{- end }}
		{{- $argNames := $argNameList.Strings }}
		{{- $argTypes := $argTypeList.Strings }}

	{{/* =========================== */}}
	{{/*      unique index code      */}}
	{{/* =========================== */}}

// {{ $structName }}By{{ $index.PascalName }} query {{ printf "%+q" $tableName }} table by {{ if $index.Primary }}primary key{{ else }}unique key {{ printf "%+q" $index.Name }}{{ end }}.
// Return nil if error occurred or there is not row found.
func {{ $structName }}By{{ $index.PascalName }}(ctx_ {{ $ctx }}.Context, db_ DBer{{ range $j, $argName := $argNames }}, {{ $argName }} {{ index $argTypes $j }}{{ end }}) (*{{ $structName }}, error) {

	sql_ := {{ $sqlx }}.Rebind(BindType, "SELECT {{ join (columnNames $cols) ", " }} " +
		"FROM {{ $tableName }} " +
		"WHERE {{ range $j, $col := $indexCols }}{{ if ne $j 0 }}AND {{ end }}{{ $col.Name }}=? {{ end }}")

	row_ := db_.QueryRowContext(ctx_, sql_{{ range $j, $argName := $argNames }}, {{ $argName }}{{ end }})

	entry_ := new({{ $structName }})
	if err_ := row_.Scan({{ range $j, $field := $structFieldNames }}{{ if ne $j 0 }}, {{ end }}&entry_.{{ $field }}{{ end }}); err_ != nil {
		if err_ == {{ $sql }}.ErrNoRows {
			return nil, nil
		}
		return nil, err_
	}

	return entry_, nil
}

	{{- end }}
{{- end }}

`)

}
<|endoftext|>


<commit_before>package user

import (
	"bytes"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/eirka/eirka-libs/config"
	"github.com/eirka/eirka-libs/db"

	local "github.com/eirka/eirka-post/config"
)

func init() {

	// Database connection settings
	dbase := db.Database{

		User:           local.Settings.Database.User,
		Password:       local.Settings.Database.Password,
		Proto:          local.Settings.Database.Proto,
		Host:           local.Settings.Database.Host,
		Database:       local.Settings.Database.Database,
		MaxIdle:        local.Settings.Database.MaxIdle,
		MaxConnections: local.Settings.Database.MaxConnections,
	}

	// Set up DB connection
	dbase.NewDb()

	// Get limits and stuff from database
	config.GetDatabaseSettings()
}

func performRequest(r http.Handler, method, path string) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, nil)
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func performJwtHeaderRequest(r http.Handler, method, path, token string) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, nil)
	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func performJwtFormRequest(r http.Handler, method, path, token string) *httptest.ResponseRecorder {
	var b bytes.Buffer

	mw := multipart.NewWriter(&b)
	mw.WriteField("access_token", token)
	mw.Close()

	req, _ := http.NewRequest(method, path, &b)
	req.Header.Set("Content-Type", mw.FormDataContentType())
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func TestAuthSecret(t *testing.T) {

	Secret = ""

	gin.SetMode(gin.ReleaseMode)

	router := gin.New()

	// route is open
	router.Use(Auth(false))

	router.GET("/", func(c *gin.Context) {
		c.String(200, "OK")
		return
	})

	first := performRequest(router, "GET", "/")

	assert.Equal(t, first.Code, 500, "HTTP request code should match")

	Secret = "secret"

	second := performRequest(router, "GET", "/")

	assert.Equal(t, second.Code, 200, "HTTP request code should match")

	user := DefaultUser()
	user.SetId(2)
	user.SetAuthenticated()

	user.hash, err = HashPassword("testpassword")
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotNil(t, user.hash, "password should be returned")
	}

	assert.True(t, user.ComparePassword("testpassword"), "Password should validate")

	badtoken, err := user.CreateToken()
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotEmpty(t, badtoken, "token should be returned")
	}

	third := performJwtHeaderRequest(router, "GET", "/", goodtoken)

	assert.Equal(t, third.Code, 200, "HTTP request code should match")

}

func TestAuthHeaderToken(t *testing.T) {

	var err error

	Secret = "secret"

	gin.SetMode(gin.ReleaseMode)

	router := gin.New()

	// route is protected
	router.Use(Auth(true))

	router.GET("/", func(c *gin.Context) {
		c.String(200, "OK")
		return
	})

	first := performRequest(router, "GET", "/")

	assert.Equal(t, first.Code, 401, "HTTP request code should match")

	user := DefaultUser()
	user.SetId(2)
	user.SetAuthenticated()

	user.hash, err = HashPassword("testpassword")
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotNil(t, user.hash, "password should be returned")
	}

	assert.True(t, user.ComparePassword("testpassword"), "Password should validate")

	badtoken, err := user.CreateToken()
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotEmpty(t, badtoken, "token should be returned")
	}

	Secret = "changed"

	second := performJwtHeaderRequest(router, "GET", "/", badtoken)

	assert.Equal(t, second.Code, 401, "HTTP request code should match")

	goodtoken, err := user.CreateToken()
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotEmpty(t, goodtoken, "token should be returned")
	}

	third := performJwtHeaderRequest(router, "GET", "/", goodtoken)

	assert.Equal(t, third.Code, 200, "HTTP request code should match")

}

func TestAuthFormToken(t *testing.T) {

	var err error

	Secret = "secret"

	gin.SetMode(gin.ReleaseMode)

	router := gin.New()

	// route is protected
	router.Use(Auth(true))

	router.GET("/", func(c *gin.Context) {
		c.String(200, "OK")
		return
	})

	user := DefaultUser()
	user.SetId(2)
	user.SetAuthenticated()

	user.hash, err = HashPassword("testpassword")
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotNil(t, user.hash, "password should be returned")
	}

	assert.True(t, user.ComparePassword("testpassword"), "Password should validate")

	badtoken, err := user.CreateToken()
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotEmpty(t, badtoken, "token should be returned")
	}

	Secret = "changed"

	first := performJwtFormRequest(router, "GET", "/", badtoken)

	assert.Equal(t, first.Code, 401, "HTTP request code should match")

	goodtoken, err := user.CreateToken()
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotEmpty(t, goodtoken, "token should be returned")
	}

	second := performJwtFormRequest(router, "GET", "/", goodtoken)

	assert.Equal(t, second.Code, 200, "HTTP request code should match")

}
<commit_msg>enhance auth test to check that logged in users can see anon routes<commit_after>package user

import (
	"bytes"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/eirka/eirka-libs/config"
	"github.com/eirka/eirka-libs/db"

	local "github.com/eirka/eirka-post/config"
)

func init() {

	// Database connection settings
	dbase := db.Database{

		User:           local.Settings.Database.User,
		Password:       local.Settings.Database.Password,
		Proto:          local.Settings.Database.Proto,
		Host:           local.Settings.Database.Host,
		Database:       local.Settings.Database.Database,
		MaxIdle:        local.Settings.Database.MaxIdle,
		MaxConnections: local.Settings.Database.MaxConnections,
	}

	// Set up DB connection
	dbase.NewDb()

	// Get limits and stuff from database
	config.GetDatabaseSettings()
}

func performRequest(r http.Handler, method, path string) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, nil)
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func performJwtHeaderRequest(r http.Handler, method, path, token string) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, nil)
	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func performJwtFormRequest(r http.Handler, method, path, token string) *httptest.ResponseRecorder {
	var b bytes.Buffer

	mw := multipart.NewWriter(&b)
	mw.WriteField("access_token", token)
	mw.Close()

	req, _ := http.NewRequest(method, path, &b)
	req.Header.Set("Content-Type", mw.FormDataContentType())
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func TestAuthSecret(t *testing.T) {

	Secret = ""

	gin.SetMode(gin.ReleaseMode)

	router := gin.New()

	// route is open
	router.Use(Auth(false))

	router.GET("/", func(c *gin.Context) {
		c.String(200, "OK")
		return
	})

	first := performRequest(router, "GET", "/")

	assert.Equal(t, first.Code, 500, "HTTP request code should match")

	Secret = "secret"

	second := performRequest(router, "GET", "/")

	assert.Equal(t, second.Code, 200, "HTTP request code should match")

	user := DefaultUser()
	user.SetId(2)
	user.SetAuthenticated()

	user.hash, err := HashPassword("testpassword")
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotNil(t, user.hash, "password should be returned")
	}

	assert.True(t, user.ComparePassword("testpassword"), "Password should validate")

	token, err := user.CreateToken()
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotEmpty(t, badtoken, "token should be returned")
	}

	third := performJwtHeaderRequest(router, "GET", "/", token)

	assert.Equal(t, third.Code, 200, "HTTP request code should match")

}

func TestAuthHeaderToken(t *testing.T) {

	var err error

	Secret = "secret"

	gin.SetMode(gin.ReleaseMode)

	router := gin.New()

	// route is protected
	router.Use(Auth(true))

	router.GET("/", func(c *gin.Context) {
		c.String(200, "OK")
		return
	})

	first := performRequest(router, "GET", "/")

	assert.Equal(t, first.Code, 401, "HTTP request code should match")

	user := DefaultUser()
	user.SetId(2)
	user.SetAuthenticated()

	user.hash, err = HashPassword("testpassword")
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotNil(t, user.hash, "password should be returned")
	}

	assert.True(t, user.ComparePassword("testpassword"), "Password should validate")

	badtoken, err := user.CreateToken()
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotEmpty(t, badtoken, "token should be returned")
	}

	Secret = "changed"

	second := performJwtHeaderRequest(router, "GET", "/", badtoken)

	assert.Equal(t, second.Code, 401, "HTTP request code should match")

	goodtoken, err := user.CreateToken()
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotEmpty(t, goodtoken, "token should be returned")
	}

	third := performJwtHeaderRequest(router, "GET", "/", goodtoken)

	assert.Equal(t, third.Code, 200, "HTTP request code should match")

}

func TestAuthFormToken(t *testing.T) {

	var err error

	Secret = "secret"

	gin.SetMode(gin.ReleaseMode)

	router := gin.New()

	// route is protected
	router.Use(Auth(true))

	router.GET("/", func(c *gin.Context) {
		c.String(200, "OK")
		return
	})

	user := DefaultUser()
	user.SetId(2)
	user.SetAuthenticated()

	user.hash, err = HashPassword("testpassword")
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotNil(t, user.hash, "password should be returned")
	}

	assert.True(t, user.ComparePassword("testpassword"), "Password should validate")

	badtoken, err := user.CreateToken()
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotEmpty(t, badtoken, "token should be returned")
	}

	Secret = "changed"

	first := performJwtFormRequest(router, "GET", "/", badtoken)

	assert.Equal(t, first.Code, 401, "HTTP request code should match")

	goodtoken, err := user.CreateToken()
	if assert.NoError(t, err, "An error was not expected") {
		assert.NotEmpty(t, goodtoken, "token should be returned")
	}

	second := performJwtFormRequest(router, "GET", "/", goodtoken)

	assert.Equal(t, second.Code, 200, "HTTP request code should match")

}
<|endoftext|>