<commit_before>package hhfrag

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"text/tabwriter"

	"github.com/BurntSushi/bcbgo/apps/hhsuite"
	"github.com/BurntSushi/bcbgo/io/hhm"
	"github.com/BurntSushi/bcbgo/io/hhr"
	"github.com/BurntSushi/bcbgo/io/pdb"
	"github.com/BurntSushi/bcbgo/seq"
)

type PDBDatabase hhsuite.Database

func (db PDBDatabase) HHsuite() hhsuite.Database {
	resolved := hhsuite.Database(db).Resolve()
	dbName := path.Base(resolved)
	return hhsuite.Database(path.Join(resolved, dbName))
}

func (db PDBDatabase) PDB() string {
	resolved := hhsuite.Database(db).Resolve()
	return path.Join(resolved, "pdb")
}

type Fragments struct {
	Frags      []Fragment
	Start, End int
}

// better returns true if f1 is 'better' than f2. Otherwise false.
func (f1 Fragments) better(f2 Fragments) bool {
	return len(f1.Frags) >= len(f2.Frags)
}

func (frags Fragments) Write(w io.Writer) {
	tabw := tabwriter.NewWriter(w, 0, 4, 4, ' ', 0)
	fmt.Fprintln(tabw, "Hit\tQuery\tTemplate\tProb\tCorrupt")
	for _, frag := range frags.Frags {
		var corruptStr string
		if frag.IsCorrupt() {
			corruptStr = "\tcorrupt"
		}
		fmt.Fprintf(tabw, "%s\t(%d-%d)\t(%d-%d)\t%f%s\n",
			frag.Template.Name,
			frag.Hit.QueryStart, frag.Hit.QueryEnd,
			frag.Hit.TemplateStart, frag.Hit.TemplateEnd,
			frag.Hit.Prob,
			corruptStr)
	}
	tabw.Flush()
}

func FindFragments(pdbDb PDBDatabase, blits bool,
	queryHHM *hhm.HHM, qs seq.Sequence, start, end int) (*Fragments, error) {

	pre := fmt.Sprintf("bcbgo-hhfrag-hhm-%d-%d_", start, end)
	hhmFile, err := ioutil.TempFile("", pre)
	if err != nil {
		return nil, err
	}
	defer os.Remove(hhmFile.Name())
	hhmName := hhmFile.Name()

	if err := hhm.Write(hhmFile, queryHHM.Slice(start, end)); err != nil {
		return nil, err
	}

	var results *hhr.HHR
	if blits {
		conf := hhsuite.HHBlitsDefault
		conf.CPUs = 1
		results, err = conf.Run(pdbDb.HHsuite(), hhmName)
	} else {
		conf := hhsuite.HHSearchDefault
		conf.CPUs = 1
		results, err = conf.Run(pdbDb.HHsuite(), hhmName)
	}
	if err != nil {
		return nil, err
	}

	frags := make([]Fragment, len(results.Hits))
	for i, hit := range results.Hits {
		hit.QueryStart += start
		hit.QueryEnd += start
		frag, err := NewFragment(pdbDb, qs, hit)
		if err != nil {
			return nil, err
		}
		frags[i] = frag
	}
	return &Fragments{
		Frags: frags,
		Start: start,
		End:   end,
	}, nil
}

// An HHfrag Fragment corresponds to a match between a portion of a query
// HMM and a portion of a template HMM. The former is represented as a slice
// of a regular sequence, where the latter is represented as an hhsuite hit
// and a list of alpha-carbon atoms corresponding to the matched region.
type Fragment struct {
	Query    seq.Sequence
	Template seq.Sequence
	Hit      hhr.Hit
	CaAtoms  []pdb.Coords
}

// IsCorrupt returns true when a particular fragment could not be paired
// with alpha-carbon positions for every residue in the template strand.
// (This problem stems from the fact that we use SEQRES records for sequence
// information, but not all residues in SEQRES have alpha-carbon ATOM records
// associated with them.)
func (frag Fragment) IsCorrupt() bool {
	return frag.CaAtoms == nil
}

// NewFragment constructs a new fragment from a full query sequence and the
// hit from the HHR file.
//
// Since NewFragment requires access to the raw PDB alpha-carbon atoms (and
// the sequence) of the template hit, you'll also need to pass a path to the
// PDB database. (Which is a directory containing a flat list of all
// PDB files used to construct the corresponding hhblits database.) This
// database is usually located inside the 'pdb' directory contained in the
// corresponding hhsuite database. i.e., $HHLIB/data/pdb-select25/pdb
func NewFragment(
	pdbDb PDBDatabase, qs seq.Sequence, hit hhr.Hit) (Fragment, error) {

	pdbName := getTemplatePdbName(hit.Name)
	pdbEntry, err := pdb.ReadPDB(path.Join(
		pdbDb.PDB(), fmt.Sprintf("%s.pdb", pdbName)))
	if err != nil {
		return Fragment{}, err
	}

	// Load in the sequence from the PDB file using the SEQRES residues.
	ts, te := hit.TemplateStart, hit.TemplateEnd
	chain := pdbEntry.OneChain()
	tseq := seq.Sequence{
		Name:     pdbName,
		Residues: make([]seq.Residue, te-ts+1),
	}

	// We copy here to avoid pinning pdb.Entry objects.
	copy(tseq.Residues, chain.Sequence[ts-1:te])

	frag := Fragment{
		Query:    qs.Slice(hit.QueryStart-1, hit.QueryEnd),
		Template: tseq,
		Hit:      hit,
		CaAtoms:  nil,
	}

	// We designate "corrupt" if the query/template hit regions are of
	// different length. i.e., we don't allow gaps (yet).
	// BUG(burntsushi): Fragments with gaps are marked as corrupt.
	if hit.QueryEnd-hit.QueryStart != hit.TemplateEnd-hit.TemplateStart {
		return frag, nil
	}

	// We also designate "corrupt" if there are any gaps in our alpha-carbon
	// atom list.
	atoms := chain.SequenceCaAtomSlice(ts-1, te)
	if atoms == nil {
		return frag, nil
	}

	// One again, we copy to avoid pinning memory.
	frag.CaAtoms = make([]pdb.Coords, len(atoms))
	copy(frag.CaAtoms, atoms)

	return frag, nil
}

func getTemplatePdbName(hitName string) string {
	return strings.SplitN(strings.TrimSpace(hitName), " ", 2)[0]
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}
<commit_msg>Split up gapped fragments.<commit_after>package hhfrag

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"text/tabwriter"

	"github.com/BurntSushi/bcbgo/apps/hhsuite"
	"github.com/BurntSushi/bcbgo/io/hhm"
	"github.com/BurntSushi/bcbgo/io/hhr"
	"github.com/BurntSushi/bcbgo/io/pdb"
	"github.com/BurntSushi/bcbgo/seq"
)

type PDBDatabase hhsuite.Database

func (db PDBDatabase) HHsuite() hhsuite.Database {
	resolved := hhsuite.Database(db).Resolve()
	dbName := path.Base(resolved)
	return hhsuite.Database(path.Join(resolved, dbName))
}

func (db PDBDatabase) PDB() string {
	resolved := hhsuite.Database(db).Resolve()
	return path.Join(resolved, "pdb")
}

type Fragments struct {
	Frags      []Fragment
	Start, End int
}

// better returns true if f1 is 'better' than f2. Otherwise false.
func (f1 Fragments) better(f2 Fragments) bool {
	return len(f1.Frags) >= len(f2.Frags)
}

func (frags Fragments) Write(w io.Writer) {
	tabw := tabwriter.NewWriter(w, 0, 4, 4, ' ', 0)
	fmt.Fprintln(tabw, "Hit\tQuery\tTemplate\tProb\tCorrupt")
	for _, frag := range frags.Frags {
		var corruptStr string
		if frag.IsCorrupt() {
			corruptStr = "\tcorrupt"
		}
		fmt.Fprintf(tabw, "%s\t(%d-%d)\t(%d-%d)\t%f%s\n",
			frag.Template.Name,
			frag.Hit.QueryStart, frag.Hit.QueryEnd,
			frag.Hit.TemplateStart, frag.Hit.TemplateEnd,
			frag.Hit.Prob,
			corruptStr)
	}
	tabw.Flush()
}

func FindFragments(pdbDb PDBDatabase, blits bool,
	queryHHM *hhm.HHM, qs seq.Sequence, start, end int) (*Fragments, error) {

	pre := fmt.Sprintf("bcbgo-hhfrag-hhm-%d-%d_", start, end)
	hhmFile, err := ioutil.TempFile("", pre)
	if err != nil {
		return nil, err
	}
	defer os.Remove(hhmFile.Name())
	hhmName := hhmFile.Name()

	if err := hhm.Write(hhmFile, queryHHM.Slice(start, end)); err != nil {
		return nil, err
	}

	var results *hhr.HHR
	if blits {
		conf := hhsuite.HHBlitsDefault
		conf.CPUs = 1
		results, err = conf.Run(pdbDb.HHsuite(), hhmName)
	} else {
		conf := hhsuite.HHSearchDefault
		conf.CPUs = 1
		results, err = conf.Run(pdbDb.HHsuite(), hhmName)
	}
	if err != nil {
		return nil, err
	}

	frags := make([]Fragment, 0, len(results.Hits))
	for _, hit := range results.Hits {
		hit.QueryStart += start
		hit.QueryEnd += start
		for _, splitted := range splitHit(hit) {
			frag, err := NewFragment(pdbDb, qs, splitted)
			if err != nil {
				return nil, err
			}
			frags = append(frags, frag)
		}
	}
	return &Fragments{
		Frags: frags,
		Start: start,
		End:   end,
	}, nil
}

func splitHit(hit hhr.Hit) []hhr.Hit {
	splitted := make([]hhr.Hit, 0)
	chunks := 0
	start := 0
	for i, r := range hit.Aligned.TSeq {
		if r == '-' {
			// Skip if this is the first residue or last residue was '-',
			// since we haven't accumulated anything.
			if i == 0 || hit.Aligned.TSeq[i-1] == '-' {
				start++
				continue
			}

			splitted = append(splitted, splitAt(hit, chunks, start, i))
			chunks++
			start = i + 1
		}
	}
	if start < len(hit.Aligned.TSeq) {
		piece := splitAt(hit, chunks, start, len(hit.Aligned.TSeq))
		splitted = append(splitted, piece)
	}
	return splitted
}

func splitAt(hit hhr.Hit, chunkNum, start, end int) hhr.Hit {
	cpy := hit

	cpy.Chunk = chunkNum
	cpy.NumAlignedCols = end - start
	cpy.QueryStart = cpy.QueryStart + start
	cpy.QueryEnd = cpy.QueryStart + cpy.NumAlignedCols - 1
	cpy.TemplateStart = cpy.TemplateStart + start
	cpy.TemplateEnd = cpy.TemplateStart + cpy.NumAlignedCols - 1
	cpy.NumTemplateCols = cpy.NumAlignedCols

	cpy.Aligned.QSeq = cpy.Aligned.QSeq[start:end]
	cpy.Aligned.QConsensus = cpy.Aligned.QConsensus[start:end]
	if len(cpy.Aligned.QDssp) > 0 {
		cpy.Aligned.QDssp = cpy.Aligned.QDssp[start:end]
	}
	if len(cpy.Aligned.QPred) > 0 {
		cpy.Aligned.QPred = cpy.Aligned.QPred[start:end]
	}
	if len(cpy.Aligned.QConf) > 0 {
		cpy.Aligned.QConf = cpy.Aligned.QConf[start:end]
	}

	cpy.Aligned.TSeq = cpy.Aligned.TSeq[start:end]
	cpy.Aligned.TConsensus = cpy.Aligned.TConsensus[start:end]
	if len(cpy.Aligned.TDssp) > 0 {
		cpy.Aligned.TDssp = cpy.Aligned.TDssp[start:end]
	}
	if len(cpy.Aligned.TPred) > 0 {
		cpy.Aligned.TPred = cpy.Aligned.TPred[start:end]
	}
	if len(cpy.Aligned.TConf) > 0 {
		cpy.Aligned.TConf = cpy.Aligned.TConf[start:end]
	}

	return cpy
}

// An HHfrag Fragment corresponds to a match between a portion of a query
// HMM and a portion of a template HMM. The former is represented as a slice
// of a regular sequence, where the latter is represented as an hhsuite hit
// and a list of alpha-carbon atoms corresponding to the matched region.
type Fragment struct {
	Query    seq.Sequence
	Template seq.Sequence
	Hit      hhr.Hit
	CaAtoms  []pdb.Coords
}

// IsCorrupt returns true when a particular fragment could not be paired
// with alpha-carbon positions for every residue in the template strand.
// (This problem stems from the fact that we use SEQRES records for sequence
// information, but not all residues in SEQRES have alpha-carbon ATOM records
// associated with them.)
func (frag Fragment) IsCorrupt() bool {
	return frag.CaAtoms == nil
}

// NewFragment constructs a new fragment from a full query sequence and the
// hit from the HHR file.
//
// Since NewFragment requires access to the raw PDB alpha-carbon atoms (and
// the sequence) of the template hit, you'll also need to pass a path to the
// PDB database. (Which is a directory containing a flat list of all
// PDB files used to construct the corresponding hhblits database.) This
// database is usually located inside the 'pdb' directory contained in the
// corresponding hhsuite database. i.e., $HHLIB/data/pdb-select25/pdb
func NewFragment(
	pdbDb PDBDatabase, qs seq.Sequence, hit hhr.Hit) (Fragment, error) {

	pdbName := getTemplatePdbName(hit.Name)
	pdbEntry, err := pdb.ReadPDB(path.Join(
		pdbDb.PDB(), fmt.Sprintf("%s.pdb", pdbName)))
	if err != nil {
		return Fragment{}, err
	}

	// Load in the sequence from the PDB file using the SEQRES residues.
	ts, te := hit.TemplateStart, hit.TemplateEnd
	chain := pdbEntry.OneChain()
	tseq := seq.Sequence{
		Name:     pdbName,
		Residues: make([]seq.Residue, te-ts+1),
	}

	// We copy here to avoid pinning pdb.Entry objects.
	copy(tseq.Residues, chain.Sequence[ts-1:te])

	frag := Fragment{
		Query:    qs.Slice(hit.QueryStart-1, hit.QueryEnd),
		Template: tseq,
		Hit:      hit,
		CaAtoms:  nil,
	}

	// We designate "corrupt" if the query/template hit regions are of
	// different length. i.e., we don't allow gaps (yet).
	// BUG(burntsushi): Fragments with gaps are marked as corrupt.
	if hit.QueryEnd-hit.QueryStart != hit.TemplateEnd-hit.TemplateStart {
		return frag, nil
	}

	// We also designate "corrupt" if there are any gaps in our alpha-carbon
	// atom list.
	atoms := chain.SequenceCaAtomSlice(ts-1, te)
	if atoms == nil {
		return frag, nil
	}

	// One again, we copy to avoid pinning memory.
	frag.CaAtoms = make([]pdb.Coords, len(atoms))
	copy(frag.CaAtoms, atoms)

	return frag, nil
}

func getTemplatePdbName(hitName string) string {
	return strings.SplitN(strings.TrimSpace(hitName), " ", 2)[0]
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"bufio"
	"fmt"
	"io"
	"sync/atomic"
	"runtime"
//"github.com/surgemq/message"
	"encoding/binary"
)

var (
	bufcnt int64
	DefaultBufferSize int64

	DeviceInBufferSize int64
	DeviceOutBufferSize int64

	MasterInBufferSize int64
	MasterOutBufferSize int64
)

const (
	smallReadBlockSize = 512
	defaultReadBlockSize = 8192
	defaultWriteBlockSize = 8192
)

/**
2016.03.03 修改
bingbuffer结构体
 */
type buffer struct {
	readIndex  int64        //读序号
	writeIndex int64        //写序号
	ringBuffer []*ByteArray //环形buffer指针数组
	bufferSize int64        //初始化环形buffer指针数组大小
	mask       int64        //掩码：bufferSize-1
	done       int64        //是否完成
}

type ByteArray struct {
	bArray []byte
}

func (this *buffer)ReadCommit(index int64) {
	this.ringBuffer[index] = nil
}

/**
2016.03.03 添加
初始化ringbuffer
参数bufferSize：初始化环形buffer指针数组大小
 */
func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}
	if size == 0 {
		size = DefaultBufferSize
	}
	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		readIndex: int64(0), //读序号
		writeIndex: int64(0), //写序号
		ringBuffer: make([]*ByteArray, size), //环形buffer指针数组
		bufferSize: size, //初始化环形buffer指针数组大小
		mask:size - 1,
	}, nil
}

/**
2016.03.03 添加
获取当前读序号
 */
func (this *buffer)GetCurrentReadIndex() (int64) {
	return atomic.LoadInt64(&this.readIndex)
}
/**
2016.03.03 添加
获取当前写序号
 */
func (this *buffer)GetCurrentWriteIndex() (int64) {
	return atomic.LoadInt64(&this.writeIndex)
}

/**
2016.03.03 添加
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
 */
func (this *buffer)ReadBuffer() ([]byte, int64, bool) {

	readIndex := atomic.LoadInt64(&this.readIndex)
	//writeIndex := atomic.LoadInt64(&this.writeIndex)
	//switch  {
	//case readIndex >= writeIndex:
	//	return nil, -1, false
	//case writeIndex - readIndex > this.bufferSize:
	//	return nil, -1, false
	//default:
	index := readIndex & this.mask

	p_ := this.ringBuffer[index]
	//this.ringBuffer[index] = nil
	if p_ == nil {
		return nil, -1, false
	}
	p := p_.bArray

	if p == nil {
		return nil, -1, false
	}
	atomic.AddInt64(&this.readIndex, 1)
	return p, index, true
	//}
	//return nil, -1, false
}

/**
2016.03.03 添加
写入ringbuffer指针，以及将写序号加1
 */
func (this *buffer)WriteBuffer(in []byte) (bool) {

	//readIndex := atomic.LoadInt64(&this.readIndex)
	writeIndex := atomic.LoadInt64(&this.writeIndex)
	//switch  {
	//case writeIndex - readIndex < 0:
	//	return false
	//default:
	index := writeIndex & this.mask
	if this.ringBuffer[index] == nil {
		this.ringBuffer[index] = &ByteArray{bArray:in}
		atomic.AddInt64(&this.writeIndex, 1)
		return true
	}else {
		return false
	}
	//}
}

/**
2016.03.03 修改
完成
 */
func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)
	return nil
}
/*

/**
2016.03.03 修改
向ringbuffer中写数据（从connection的中向ringbuffer中写）--生产者
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()
	for {
		total := int64(0)
		if this.isDone() {
			return total, io.EOF
		}
		b := make([]byte, 5)
		n, err := r.Read(b[0:1])

		if n > 0 {
			total += int64(n)
			if err != nil {
				return total, err
			}
		}

		/**************************/
		cnt := 1

		// Let's read enough bytes to get the message header (msg type, remaining length)
		for {
			// If we have read 5 bytes and still not done, then there's a problem.
			if cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}

			// Peek cnt bytes from the input buffer.
			_, err := r.Read(b[cnt:cnt + 1])
			if err != nil {
				return 0, err
			}

			// If we got enough bytes, then check the last byte to see if the continuation
			// bit is set. If so, increment cnt and continue peeking
			if b[cnt] >= 0x80 {
				cnt++
			} else {
				break
			}
		}

		// Get the remaining length of the message
		remlen, _ := binary.Uvarint(b[1:])

		// Total message length is remlen + 1 (msg type) + m (remlen bytes)
		len := int64(len(b))
		remlen_ := int64(remlen)
		total = remlen_ + int64(len)

		//mtype := message.MessageType(b[0] >> 4)
		/****************/
		//var msg message.Message
		//
		//msg, err = mtype.New()
		//if err != nil {
		//	return 0, err
		//}
		b_ := make([]byte, remlen_)
		_, err = r.Read(b_[0:])
		if err != nil {
			return 0, err
		}

		b = append(b, b_...)
		//n, err = msg.Decode(b)
		//if err != nil {
		//	return 0, err
		//}

		/*************************/

		for !this.WriteBuffer(b) {
			runtime.Gosched()
		}

		return total, nil
	}
}

/**
2016.03.03 修改
 */
func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()
	total := int64(0)
	for {
		if this.isDone() {
			return total, io.EOF
		}
		p, index, ok := this.ReadBuffer()
		if !ok {
			runtime.Gosched()
			continue
		}

		Log.Debugc(func() string {
			return fmt.Sprintf("defer this.ReadCommit(%s)", index)
		})
		Log.Debugc(func() string {
			return fmt.Sprintf("WriteTo函数》》读取*p：" + string(p))
		})

		Log.Debugc(func() string {
			return fmt.Sprintf(" WriteTo(w io.Writer)(7)")
		})
		//
		//Log.Errorc(func() string {
		//	return fmt.Sprintf("msg::" + msg.Name())
		//})
		//
		//p := make([]byte, msg.Len())
		//_, err := msg.Encode(p)
		//if err != nil {
		//	Log.Errorc(func() string {
		//		return fmt.Sprintf("msg.Encode(p)")
		//	})
		//	return total, io.EOF
		//}
		// There's some data, let's process it first
		if len(p) > 0 {
			n, err := w.Write(p)
			total += int64(n)
			Log.Debugc(func() string {
				return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)
			})

			if err != nil {
				Log.Errorc(func() string {
					return fmt.Sprintf("w.Write(p) error")
				})
				return total, err
			}
		}
		this.ReadCommit(index)
		return total, nil
	}
}

/**
2016.03.03 修改
*/
func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n & (n - 1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}
<commit_msg>修改buffer读写限制<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"bufio"
	"fmt"
	"io"
	"sync/atomic"
	"runtime"
//"github.com/surgemq/message"
	"encoding/binary"
)

var (
	bufcnt int64
	DefaultBufferSize int64

	DeviceInBufferSize int64
	DeviceOutBufferSize int64

	MasterInBufferSize int64
	MasterOutBufferSize int64
)

const (
	smallReadBlockSize = 512
	defaultReadBlockSize = 8192
	defaultWriteBlockSize = 8192
)

/**
2016.03.03 修改
bingbuffer结构体
 */
type buffer struct {
	readIndex  int64        //读序号
	writeIndex int64        //写序号
	ringBuffer []*ByteArray //环形buffer指针数组
	bufferSize int64        //初始化环形buffer指针数组大小
	mask       int64        //掩码：bufferSize-1
	done       int64        //是否完成
}

type ByteArray struct {
	bArray []byte
}

func (this *buffer)ReadCommit(index int64) {
	this.ringBuffer[index] = nil
}

/**
2016.03.03 添加
初始化ringbuffer
参数bufferSize：初始化环形buffer指针数组大小
 */
func newBuffer(size int64) (*buffer, error) {
	if size < 0 {
		return nil, bufio.ErrNegativeCount
	}
	if size == 0 {
		size = DefaultBufferSize
	}
	if !powerOfTwo64(size) {
		fmt.Printf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
		return nil, fmt.Errorf("Size must be power of two. Try %d.", roundUpPowerOfTwo64(size))
	}

	return &buffer{
		readIndex: int64(0), //读序号
		writeIndex: int64(0), //写序号
		ringBuffer: make([]*ByteArray, size), //环形buffer指针数组
		bufferSize: size, //初始化环形buffer指针数组大小
		mask:size - 1,
	}, nil
}

/**
2016.03.03 添加
获取当前读序号
 */
func (this *buffer)GetCurrentReadIndex() (int64) {
	return atomic.LoadInt64(&this.readIndex)
}
/**
2016.03.03 添加
获取当前写序号
 */
func (this *buffer)GetCurrentWriteIndex() (int64) {
	return atomic.LoadInt64(&this.writeIndex)
}

/**
2016.03.03 添加
读取ringbuffer指定的buffer指针，返回该指针并清空ringbuffer该位置存在的指针内容，以及将读序号加1
 */
func (this *buffer)ReadBuffer() ([]byte, int64, bool) {

	readIndex := atomic.LoadInt64(&this.readIndex)
	//writeIndex := atomic.LoadInt64(&this.writeIndex)
	//switch  {
	//case readIndex >= writeIndex:
	//	return nil, -1, false
	//case writeIndex - readIndex > this.bufferSize:
	//	return nil, -1, false
	//default:

	atomic.AddInt64(&this.readIndex, 1)
	index := readIndex & this.mask

	p_ := this.ringBuffer[index]
	//this.ringBuffer[index] = nil
	if p_ == nil {
		return nil, -1, false
	}
	p := p_.bArray

	if p == nil {
		return nil, -1, false
	}

	return p, index, true
	//}
	//return nil, -1, false
}

/**
2016.03.03 添加
写入ringbuffer指针，以及将写序号加1
 */
func (this *buffer)WriteBuffer(in []byte) (bool) {

	//readIndex := atomic.LoadInt64(&this.readIndex)
	writeIndex := atomic.LoadInt64(&this.writeIndex)
	//switch  {
	//case writeIndex - readIndex < 0:
	//	return false
	//default:
	index := writeIndex & this.mask
	if this.ringBuffer[index] == nil {
		this.ringBuffer[index] = &ByteArray{bArray:in}
		atomic.AddInt64(&this.writeIndex, 1)
		return true
	}else {
		return false
	}
	//}
}

/**
2016.03.03 修改
完成
 */
func (this *buffer) Close() error {
	atomic.StoreInt64(&this.done, 1)
	return nil
}
/*

/**
2016.03.03 修改
向ringbuffer中写数据（从connection的中向ringbuffer中写）--生产者
*/
func (this *buffer) ReadFrom(r io.Reader) (int64, error) {
	defer this.Close()
	for {
		total := int64(0)
		if this.isDone() {
			return total, io.EOF
		}
		b := make([]byte, 5)
		n, err := r.Read(b[0:1])

		if n > 0 {
			total += int64(n)
			if err != nil {
				return total, err
			}
		}

		/**************************/
		cnt := 1

		// Let's read enough bytes to get the message header (msg type, remaining length)
		for {
			// If we have read 5 bytes and still not done, then there's a problem.
			if cnt > 4 {
				return 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
			}

			// Peek cnt bytes from the input buffer.
			_, err := r.Read(b[cnt:cnt + 1])
			if err != nil {
				return 0, err
			}

			// If we got enough bytes, then check the last byte to see if the continuation
			// bit is set. If so, increment cnt and continue peeking
			if b[cnt] >= 0x80 {
				cnt++
			} else {
				break
			}
		}

		// Get the remaining length of the message
		remlen, _ := binary.Uvarint(b[1:])

		// Total message length is remlen + 1 (msg type) + m (remlen bytes)
		len := int64(len(b))
		remlen_ := int64(remlen)
		total = remlen_ + int64(len)

		//mtype := message.MessageType(b[0] >> 4)
		/****************/
		//var msg message.Message
		//
		//msg, err = mtype.New()
		//if err != nil {
		//	return 0, err
		//}
		b_ := make([]byte, remlen_)
		_, err = r.Read(b_[0:])
		if err != nil {
			return 0, err
		}

		b = append(b, b_...)
		//n, err = msg.Decode(b)
		//if err != nil {
		//	return 0, err
		//}

		/*************************/

		for !this.WriteBuffer(b) {
			runtime.Gosched()
		}

		return total, nil
	}
}

/**
2016.03.03 修改
 */
func (this *buffer) WriteTo(w io.Writer) (int64, error) {
	defer this.Close()
	total := int64(0)
	for {
		if this.isDone() {
			return total, io.EOF
		}
		p, index, ok := this.ReadBuffer()
		if !ok {
			runtime.Gosched()
			continue
		}

		Log.Debugc(func() string {
			return fmt.Sprintf("defer this.ReadCommit(%s)", index)
		})
		Log.Debugc(func() string {
			return fmt.Sprintf("WriteTo函数》》读取*p：" + string(p))
		})

		Log.Debugc(func() string {
			return fmt.Sprintf(" WriteTo(w io.Writer)(7)")
		})
		//
		//Log.Errorc(func() string {
		//	return fmt.Sprintf("msg::" + msg.Name())
		//})
		//
		//p := make([]byte, msg.Len())
		//_, err := msg.Encode(p)
		//if err != nil {
		//	Log.Errorc(func() string {
		//		return fmt.Sprintf("msg.Encode(p)")
		//	})
		//	return total, io.EOF
		//}
		// There's some data, let's process it first
		if len(p) > 0 {
			n, err := w.Write(p)
			total += int64(n)
			Log.Debugc(func() string {
				return fmt.Sprintf("Wrote %d bytes, totaling %d bytes", n, total)
			})

			if err != nil {
				Log.Errorc(func() string {
					return fmt.Sprintf("w.Write(p) error")
				})
				return total, err
			}
		}
		this.ReadCommit(index)
		return total, nil
	}
}

/**
2016.03.03 修改
*/
func (this *buffer) isDone() bool {
	if atomic.LoadInt64(&this.done) == 1 {
		return true
	}

	return false
}

func powerOfTwo64(n int64) bool {
	return n != 0 && (n & (n - 1)) == 0
}

func roundUpPowerOfTwo64(n int64) int64 {
	n--
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	n |= n >> 32
	n++

	return n
}
<|endoftext|>


<commit_before>package goupnp

import (
	"fmt"

	"github.com/huin/goupnp/soap"
)

// ServiceClient is a SOAP client, root device and the service for the SOAP
// client rolled into one value. The root device and service are intended to be
// informational.
type ServiceClient struct {
	SOAPClient *soap.SOAPClient
	RootDevice *RootDevice
	Service    *Service
}

func NewServiceClients(searchTarget string) (clients []ServiceClient, errors []error, err error) {
	var maybeRootDevices []MaybeRootDevice
	if maybeRootDevices, err = DiscoverDevices(searchTarget); err != nil {
		return
	}

	clients = make([]ServiceClient, 0, len(maybeRootDevices))

	for _, maybeRootDevice := range maybeRootDevices {
		if maybeRootDevice.Err != nil {
			errors = append(errors, maybeRootDevice.Err)
			continue
		}

		device := &maybeRootDevice.Root.Device
		srvs := device.FindService(searchTarget)
		if len(srvs) == 0 {
			errors = append(errors, fmt.Errorf("goupnp: service %q not found within device %q (UDN=%q)",
				searchTarget, device.FriendlyName, device.UDN))
			continue
		}

		for _, srv := range srvs {
			clients = append(clients, ServiceClient{
				SOAPClient: srv.NewSOAPClient(),
				RootDevice: maybeRootDevice.Root,
				Service:    srv,
			})
		}
	}

	return
}

// GetServiceClient returns the ServiceClient itself. This is provided so that the
// service client attributes can be accessed via an interface method on a
// wrapping type.
func (client *ServiceClient) GetServiceClient() *ServiceClient {
	return client
}
<commit_msg>Add non-discovery ServiceClient constructor functions.<commit_after>package goupnp

import (
	"fmt"
	"net/url"

	"github.com/huin/goupnp/soap"
)

// ServiceClient is a SOAP client, root device and the service for the SOAP
// client rolled into one value. The root device, location, and service are
// intended to be informational. Location can be used to later recreate a
// ServiceClient with NewServiceClientByURL if the service is still present;
// bypassing the discovery process.
type ServiceClient struct {
	SOAPClient *soap.SOAPClient
	RootDevice *RootDevice
	Location   *url.URL
	Service    *Service
}

// NewServiceClients discovers services, and returns clients for them. err will
// report any error with the discovery process (blocking any device/service
// discovery), errors reports errors on a per-root-device basis.
func NewServiceClients(searchTarget string) (clients []ServiceClient, errors []error, err error) {
	var maybeRootDevices []MaybeRootDevice
	if maybeRootDevices, err = DiscoverDevices(searchTarget); err != nil {
		return
	}

	clients = make([]ServiceClient, 0, len(maybeRootDevices))

	for _, maybeRootDevice := range maybeRootDevices {
		if maybeRootDevice.Err != nil {
			errors = append(errors, maybeRootDevice.Err)
			continue
		}

		deviceClients, err := NewServiceClientsFromRootDevice(maybeRootDevice.Root, maybeRootDevice.Location, searchTarget)
		if err != nil {
			errors = append(errors, err)
			continue
		}
		clients = append(clients, deviceClients...)
	}

	return
}

// NewServiceClientsByURL creates client(s) for the given service URN, for a
// root device at the given URL.
func NewServiceClientsByURL(loc *url.URL, searchTarget string) ([]ServiceClient, error) {
	rootDevice, err := DeviceByURL(loc)
	if err != nil {
		return nil, err
	}
	return NewServiceClientsFromRootDevice(rootDevice, loc, searchTarget)
}

// NewServiceClientsFromDevice creates client(s) for the given service URN, for
// a given root device. The location parameter is simply assigned to the
// Location attribute of the returned ServiceClient.
func NewServiceClientsFromRootDevice(rootDevice *RootDevice, location *url.URL, searchTarget string) ([]ServiceClient, error) {
	device := &rootDevice.Device
	srvs := device.FindService(searchTarget)
	if len(srvs) == 0 {
		return nil, fmt.Errorf("goupnp: service %q not found within device %q (UDN=%q)",
			searchTarget, device.FriendlyName, device.UDN)
	}

	clients := make([]ServiceClient, 0, len(srvs))
	for _, srv := range srvs {
		clients = append(clients, ServiceClient{
			SOAPClient: srv.NewSOAPClient(),
			RootDevice: rootDevice,
			Location:   location,
			Service:    srv,
		})
	}
	return clients, nil
}

// GetServiceClient returns the ServiceClient itself. This is provided so that the
// service client attributes can be accessed via an interface method on a
// wrapping type.
func (client *ServiceClient) GetServiceClient() *ServiceClient {
	return client
}
<|endoftext|>


<commit_before>package command

import (
	"context"

	grpc "google.golang.org/grpc"
	"v2ray.com/core"
	"v2ray.com/core/common"
	"v2ray.com/core/proxy"
)

type InboundOperation interface {
	ApplyInbound(context.Context, core.InboundHandler) error
}

type OutboundOperation interface {
	ApplyOutbound(context.Context, core.OutboundHandler) error
}

func (op *AddUserOperation) ApplyInbound(ctx context.Context, handler core.InboundHandler) error {
	getInbound, ok := handler.(proxy.GetInbound)
	if !ok {
		return newError("can't get inbound proxy from handler")
	}
	p := getInbound.GetInbound()
	um, ok := p.(proxy.UserManager)
	if !ok {
		return newError("proxy is not an UserManager")
	}
	return um.AddUser(ctx, op.User)
}

func (op *AddUserOperation) ApplyOutbound(ctx context.Context, handler core.OutboundHandler) error {
	getOutbound, ok := handler.(proxy.GetOutbound)
	if !ok {
		return newError("can't get outbound proxy from handler")
	}
	p := getOutbound.GetOutbound()
	um, ok := p.(proxy.UserManager)
	if !ok {
		return newError("proxy in not an UserManager")
	}
	return um.AddUser(ctx, op.User)
}

type handlerServer struct {
	s   *core.Instance
	ihm core.InboundHandlerManager
	ohm core.OutboundHandlerManager
}

func (s *handlerServer) AddInbound(ctx context.Context, request *AddInboundRequest) (*AddInboundResponse, error) {
	rawHandler, err := s.s.CreateObject(request.Inbound)
	if err != nil {
		return nil, err
	}
	handler, ok := rawHandler.(core.InboundHandler)
	if !ok {
		return nil, newError("not an InboundHandler.")
	}
	return &AddInboundResponse{}, s.ihm.AddHandler(ctx, handler)
}

func (s *handlerServer) RemoveInbound(ctx context.Context, request *RemoveInboundRequest) (*RemoveInboundResponse, error) {
	return &RemoveInboundResponse{}, s.ihm.RemoveHandler(ctx, request.Tag)
}

func (s *handlerServer) AlterInbound(ctx context.Context, request *AlterInboundRequest) (*AlterInboundResponse, error) {
	rawOperation, err := request.Operation.GetInstance()
	if err != nil {
		return nil, newError("unknown operation").Base(err)
	}
	operation, ok := rawOperation.(InboundOperation)
	if !ok {
		return nil, newError("not an inbound operation")
	}

	handler, err := s.ihm.GetHandler(ctx, request.Tag)
	if err != nil {
		return nil, newError("failed to get handler: ", request.Tag).Base(err)
	}

	return &AlterInboundResponse{}, operation.ApplyInbound(ctx, handler)
}

func (s *handlerServer) AddOutbound(ctx context.Context, request *AddOutboundRequest) (*AddOutboundResponse, error) {
	rawHandler, err := s.s.CreateObject(request.Outbound)
	if err != nil {
		return nil, err
	}
	handler, ok := rawHandler.(core.OutboundHandler)
	if !ok {
		return nil, newError("not an OutboundHandler.")
	}
	return &AddOutboundResponse{}, s.ohm.AddHandler(ctx, handler)
}

func (s *handlerServer) RemoveOutbound(ctx context.Context, request *RemoveOutboundRequest) (*RemoveOutboundResponse, error) {
	return &RemoveOutboundResponse{}, s.ohm.RemoveHandler(ctx, request.Tag)
}

func (s *handlerServer) AlterOutbound(ctx context.Context, request *AlterOutboundRequest) (*AlterOutboundResponse, error) {
	rawOperation, err := request.Operation.GetInstance()
	if err != nil {
		return nil, newError("unknown operation").Base(err)
	}
	operation, ok := rawOperation.(OutboundOperation)
	if !ok {
		return nil, newError("not an outbound operation")
	}

	handler := s.ohm.GetHandler(request.Tag)
	return &AlterOutboundResponse{}, operation.ApplyOutbound(ctx, handler)
}

type feature struct{}

func (*feature) Start() error {
	return nil
}

func (*feature) Close() error {
	return nil
}

func init() {
	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, cfg interface{}) (interface{}, error) {
		s := core.FromContext(ctx)
		if s == nil {
			return nil, newError("V is not in context.")
		}
		s.Commander().RegisterService(func(server *grpc.Server) {
			RegisterHandlerServiceServer(server, &handlerServer{
				s:   s,
				ihm: s.InboundHandlerManager(),
				ohm: s.OutboundHandlerManager(),
			})
		})
		return &feature{}, nil
	}))
}
<commit_msg>implement remove user operation<commit_after>package command

import (
	"context"

	grpc "google.golang.org/grpc"
	"v2ray.com/core"
	"v2ray.com/core/common"
	"v2ray.com/core/proxy"
)

// InboundOperation is the interface for operations that applies to inbound handlers.
type InboundOperation interface {
	// ApplyInbound appliess this operation to the given inbound handler.
	ApplyInbound(context.Context, core.InboundHandler) error
}

// OutboundOperation is the interface for operations that applies to outbound handlers.
type OutboundOperation interface {
	// ApplyOutbound applies this operation to the given outbound handler.
	ApplyOutbound(context.Context, core.OutboundHandler) error
}

func getInbound(handler core.InboundHandler) (proxy.Inbound, error) {
	gi, ok := handler.(proxy.GetInbound)
	if !ok {
		return nil, newError("can't get inbound proxy from handler.")
	}
	return gi.GetInbound(), nil
}

// ApplyInbound implements InboundOperation.
func (op *AddUserOperation) ApplyInbound(ctx context.Context, handler core.InboundHandler) error {
	p, err := getInbound(handler)
	if err != nil {
		return err
	}
	um, ok := p.(proxy.UserManager)
	if !ok {
		return newError("proxy is not an UserManager")
	}
	return um.AddUser(ctx, op.User)
}

// ApplyInbound implements InboundOperation.
func (op *RemoveUserOperation) ApplyInbound(ctx context.Context, handler core.InboundHandler) error {
	p, err := getInbound(handler)
	if err != nil {
		return err
	}
	um, ok := p.(proxy.UserManager)
	if !ok {
		return newError("proxy is not an UserManager")
	}
	return um.RemoveUser(ctx, op.Email)
}

type handlerServer struct {
	s   *core.Instance
	ihm core.InboundHandlerManager
	ohm core.OutboundHandlerManager
}

func (s *handlerServer) AddInbound(ctx context.Context, request *AddInboundRequest) (*AddInboundResponse, error) {
	rawHandler, err := s.s.CreateObject(request.Inbound)
	if err != nil {
		return nil, err
	}
	handler, ok := rawHandler.(core.InboundHandler)
	if !ok {
		return nil, newError("not an InboundHandler.")
	}
	return &AddInboundResponse{}, s.ihm.AddHandler(ctx, handler)
}

func (s *handlerServer) RemoveInbound(ctx context.Context, request *RemoveInboundRequest) (*RemoveInboundResponse, error) {
	return &RemoveInboundResponse{}, s.ihm.RemoveHandler(ctx, request.Tag)
}

func (s *handlerServer) AlterInbound(ctx context.Context, request *AlterInboundRequest) (*AlterInboundResponse, error) {
	rawOperation, err := request.Operation.GetInstance()
	if err != nil {
		return nil, newError("unknown operation").Base(err)
	}
	operation, ok := rawOperation.(InboundOperation)
	if !ok {
		return nil, newError("not an inbound operation")
	}

	handler, err := s.ihm.GetHandler(ctx, request.Tag)
	if err != nil {
		return nil, newError("failed to get handler: ", request.Tag).Base(err)
	}

	return &AlterInboundResponse{}, operation.ApplyInbound(ctx, handler)
}

func (s *handlerServer) AddOutbound(ctx context.Context, request *AddOutboundRequest) (*AddOutboundResponse, error) {
	rawHandler, err := s.s.CreateObject(request.Outbound)
	if err != nil {
		return nil, err
	}
	handler, ok := rawHandler.(core.OutboundHandler)
	if !ok {
		return nil, newError("not an OutboundHandler.")
	}
	return &AddOutboundResponse{}, s.ohm.AddHandler(ctx, handler)
}

func (s *handlerServer) RemoveOutbound(ctx context.Context, request *RemoveOutboundRequest) (*RemoveOutboundResponse, error) {
	return &RemoveOutboundResponse{}, s.ohm.RemoveHandler(ctx, request.Tag)
}

func (s *handlerServer) AlterOutbound(ctx context.Context, request *AlterOutboundRequest) (*AlterOutboundResponse, error) {
	rawOperation, err := request.Operation.GetInstance()
	if err != nil {
		return nil, newError("unknown operation").Base(err)
	}
	operation, ok := rawOperation.(OutboundOperation)
	if !ok {
		return nil, newError("not an outbound operation")
	}

	handler := s.ohm.GetHandler(request.Tag)
	return &AlterOutboundResponse{}, operation.ApplyOutbound(ctx, handler)
}

type feature struct{}

func (*feature) Start() error {
	return nil
}

func (*feature) Close() error {
	return nil
}

func init() {
	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, cfg interface{}) (interface{}, error) {
		s := core.FromContext(ctx)
		if s == nil {
			return nil, newError("V is not in context.")
		}
		s.Commander().RegisterService(func(server *grpc.Server) {
			RegisterHandlerServiceServer(server, &handlerServer{
				s:   s,
				ihm: s.InboundHandlerManager(),
				ohm: s.OutboundHandlerManager(),
			})
		})
		return &feature{}, nil
	}))
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"time"

	"github.com/Syfaro/telegram-bot-api"
)

const (
	// VERSION - current version
	VERSION = "1.2"

	// DEFAULT_BOT_TIMEOUT - bot default timeout
	DEFAULT_BOT_TIMEOUT = 60

	// MESSAGES_QUEUE_SIZE - size of channel for bot messages
	MESSAGES_QUEUE_SIZE = 10

	// MAX_MESSAGE_LENGTH - max length of one bot message
	MAX_MESSAGE_LENGTH = 4096

	// SECONDS_FOR_AUTO_SAVE_USERS_TO_DB - save users to file every 1 min (if need)
	SECONDS_FOR_AUTO_SAVE_USERS_TO_DB = 60

	// DB_FILE_NAME - DB json name
	DB_FILE_NAME = "shell2telegram.json"
)

// Command - one user command
type Command struct {
	shellCmd    string   // shell command
	description string   // command description for list in /help (/cmd:desc="Command name")
	vars        []string // environment vars for user text, split by `/s+` to vars (/cmd:vars=SUBCOMMAND,ARGS)
}

// Commands - list of all commands
type Commands map[string]Command

// Config - config struct
type Config struct {
	token                  string   // bot token
	addExit                bool     // adding /shell2telegram exit command
	botTimeout             int      // bot timeout
	predefinedAllowedUsers []string // telegram users who are allowed to chat with the bot
	predefinedRootUsers    []string // telegram users, who confirms new users in their private chat
	allowAll               bool     // allow all user (DANGEROUS!)
	logCommands            bool     // logging all commands
	description            string   // description of bot
	persistentUsers        bool     // load/save users from file
	usersDB                string   // file for store users
}

// message types
const (
	msgIsText int8 = iota
	msgIsPhoto
)

// BotMessage - record for send via channel for send message to telegram chat
type BotMessage struct {
	chatID      int
	messageType int8
	message     string
	fileName    string
	photo       []byte
}

// ----------------------------------------------------------------------------
// get config
func getConfig() (commands Commands, appConfig Config, err error) {
	flag.StringVar(&appConfig.token, "tb-token", "", "setting bot token (or set TB_TOKEN variable)")
	flag.BoolVar(&appConfig.addExit, "add-exit", false, "adding \"/shell2telegram exit\" command for terminate bot (for roots only)")
	flag.IntVar(&appConfig.botTimeout, "timeout", DEFAULT_BOT_TIMEOUT, "setting timeout for bot")
	flag.BoolVar(&appConfig.allowAll, "allow-all", false, "allow all users (DANGEROUS!)")
	flag.BoolVar(&appConfig.logCommands, "log-commands", false, "logging all commands")
	flag.StringVar(&appConfig.description, "description", "", "setting description of bot")
	flag.BoolVar(&appConfig.persistentUsers, "persistent_users", false, "load/save users from file (default ~/.config/shell2telegram.json)")
	flag.StringVar(&appConfig.usersDB, "users_db", "", "file for store users")
	logFilename := flag.String("log", "", "log filename, default - STDOUT")
	predefinedAllowedUsers := flag.String("allow-users", "", "telegram users who are allowed to chat with the bot (\"user1,user2\")")
	predefinedRootUsers := flag.String("root-users", "", "telegram users, who confirms new users in their private chat (\"user1,user2\")")
	version := flag.Bool("version", false, "get version")

	flag.Usage = func() {
		fmt.Printf("usage: %s [options] %s\n%s\n%s\n\noptions:\n",
			os.Args[0],
			`/chat_command "shell command" /chat_command2 "shell command2"`,
			"All text after /chat_command will be sent to STDIN of shell command.",
			"If chat command is /:plain_text - get user message without any /command (for private chats only)",
		)
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.Parse()

	if *version {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	// setup log file
	if len(*logFilename) > 0 {
		fhLog, err := os.OpenFile(*logFilename, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
		if err != nil {
			log.Fatalf("error opening log file: %v", err)
		}
		log.SetOutput(fhLog)
	}

	// setup users and roots
	if *predefinedAllowedUsers != "" {
		appConfig.predefinedAllowedUsers = strings.Split(*predefinedAllowedUsers, ",")
	}
	if *predefinedRootUsers != "" {
		appConfig.predefinedRootUsers = strings.Split(*predefinedRootUsers, ",")
	}

	commands = Commands{}
	// need >= 2 arguments and count of it must be even
	args := flag.Args()
	if len(args) < 2 || len(args)%2 == 1 {
		return commands, appConfig, fmt.Errorf("error: need pairs of /chat-command and shell-command")
	}

	for i := 0; i < len(args); i += 2 {
		path, command, err := parseBotCommand(args[i], args[i+1]) // (/path, shell_command)
		if err != nil {
			return commands, appConfig, err
		}
		commands[path] = command
	}

	if appConfig.token == "" {
		if appConfig.token = os.Getenv("TB_TOKEN"); appConfig.token == "" {
			return commands, appConfig, fmt.Errorf("TB_TOKEN environment var not found. See https://core.telegram.org/bots#botfather for more information\n")
		}
	}

	return commands, appConfig, nil
}

// ----------------------------------------------------------------------------
func sendMessage(messageSignal chan<- BotMessage, chatID int, message []byte) {
	go func() {
		fileName := ""
		fileType := http.DetectContentType(message)
		switch fileType {
		case "image/png":
			fileName = "file.png"
		case "image/jpeg":
			fileName = "file.jpeg"
		case "image/gif":
			fileName = "file.gif"
		case "image/bmp":
			fileName = "file.bmp"
		default:
			fileName = "message"
		}

		if fileName == "message" {

			// is text message
			messageString := string(message)
			messagesList := []string{}

			if len(messageString) <= MAX_MESSAGE_LENGTH {
				messagesList = []string{messageString}
			} else {
				messagesList = splitStringLinesBySize(messageString, MAX_MESSAGE_LENGTH)
			}

			for _, messageChunk := range messagesList {
				messageSignal <- BotMessage{
					chatID:      chatID,
					messageType: msgIsText,
					message:     messageChunk,
				}
			}

		} else {
			// is image
			messageSignal <- BotMessage{
				chatID:      chatID,
				messageType: msgIsPhoto,
				fileName:    fileName,
				photo:       message,
			}
		}
	}()
}

// ----------------------------------------------------------------------------
func main() {
	commands, appConfig, err := getConfig()
	if err != nil {
		log.Fatal(err)
	}

	bot, err := tgbotapi.NewBotAPI(appConfig.token)
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Authorized on bot account: @%s", bot.Self.UserName)

	tgbotConfig := tgbotapi.NewUpdate(0)
	tgbotConfig.Timeout = appConfig.botTimeout
	err = bot.UpdatesChan(tgbotConfig)
	if err != nil {
		log.Fatal(err)
	}

	users := NewUsers(appConfig)
	messageSignal := make(chan BotMessage, MESSAGES_QUEUE_SIZE)
	vacuumTicker := time.Tick(SECONDS_FOR_OLD_USERS_BEFORE_VACUUM * time.Second)
	saveToBDTicker := make(<-chan time.Time)
	autoSaveOnExitSignal := make(chan os.Signal)

	if appConfig.persistentUsers {
		saveToBDTicker = time.Tick(SECONDS_FOR_AUTO_SAVE_USERS_TO_DB * time.Second)
		signal.Notify(autoSaveOnExitSignal, os.Interrupt, os.Kill)
	}
	exitSignal := make(chan struct{})

	// all /shell2telegram sub-commands handlers
	internalCommands := map[string]func(Ctx) string{
		"stat":              cmdShell2telegramStat,
		"ban":               cmdShell2telegramBan,
		"search":            cmdShell2telegramSearch,
		"desc":              cmdShell2telegramDesc,
		"rm":                cmdShell2telegramRm,
		"exit":              cmdShell2telegramExit,
		"version":           cmdShell2telegramVersion,
		"broadcast_to_root": cmdShell2telegramBroadcastToRoot,
		"message_to_user":   cmdShell2telegramMessageToUser,
	}

	doExit := false
	for !doExit {
		select {
		case telegramUpdate := <-bot.Updates:

			var messageCmd, messageArgs string
			allUserMessage := telegramUpdate.Message.Text
			if len(allUserMessage) > 0 && allUserMessage[0] == '/' {
				messageCmd, messageArgs = splitStringHalfBySpace(allUserMessage)
			} else {
				messageCmd, messageArgs = "/:plain_text", allUserMessage
			}

			allowPlainText := false
			if _, ok := commands["/:plain_text"]; ok {
				allowPlainText = true
			}

			replayMsg := ""

			if len(messageCmd) > 0 && (messageCmd != "/:plain_text" || allowPlainText) {

				users.AddNew(telegramUpdate.Message)
				userID := telegramUpdate.Message.From.ID
				allowExec := appConfig.allowAll || users.IsAuthorized(userID)

				ctx := Ctx{
					appConfig:     &appConfig,
					users:         &users,
					commands:      commands,
					userID:        userID,
					allowExec:     allowExec,
					messageCmd:    messageCmd,
					messageArgs:   messageArgs,
					messageSignal: messageSignal,
					chatID:        telegramUpdate.Message.Chat.ID,
					exitSignal:    exitSignal,
				}

				switch {
				// commands .................................
				case messageCmd == "/auth" || messageCmd == "/authroot":
					replayMsg = cmdAuth(ctx)

				case messageCmd == "/help":
					replayMsg = cmdHelp(ctx)

				case messageCmd == "/shell2telegram" && users.IsRoot(userID):
					messageSubCmd, messageArgs := splitStringHalfBySpace(messageArgs)
					ctx.messageArgs = messageArgs
					if cmdHandler, ok := internalCommands[messageSubCmd]; ok {
						replayMsg = cmdHandler(ctx)
					} else {
						replayMsg = "Sub-command not found"
					}

				case allowExec && (allowPlainText && messageCmd == "/:plain_text" || messageCmd[0] == '/'):
					_ = cmdUser(ctx)

				} // switch for commands

				if appConfig.logCommands {
					log.Printf("%s: %s", users.String(userID), allUserMessage)
				}

				sendMessage(messageSignal, telegramUpdate.Message.Chat.ID, []byte(replayMsg))
			}

		case botMessage := <-messageSignal:
			switch {
			case botMessage.messageType == msgIsText && !stringIsEmpty(botMessage.message):
				_, err = bot.SendMessage(tgbotapi.NewMessage(botMessage.chatID, botMessage.message))
			case botMessage.messageType == msgIsPhoto && len(botMessage.photo) > 0:
				bytesPhoto := tgbotapi.FileBytes{Name: botMessage.fileName, Bytes: botMessage.photo}
				_, err = bot.SendPhoto(tgbotapi.NewPhotoUpload(botMessage.chatID, bytesPhoto))
			}

			if err != nil {
				log.Print("Bot send message error: ", err)
			}

		case <-saveToBDTicker:
			users.SaveToDB(appConfig.usersDB)

		case <-vacuumTicker:
			users.ClearOldUsers()

		case <-autoSaveOnExitSignal:
			users.needSaveDB = true
			users.SaveToDB(appConfig.usersDB)
			doExit = true

		case <-exitSignal:
			if appConfig.persistentUsers {
				users.needSaveDB = true
				users.SaveToDB(appConfig.usersDB)
			}
			doExit = true
		}
	}
}
<commit_msg>Refactoring exit by system signal<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"time"

	"github.com/Syfaro/telegram-bot-api"
)

const (
	// VERSION - current version
	VERSION = "1.2"

	// DEFAULT_BOT_TIMEOUT - bot default timeout
	DEFAULT_BOT_TIMEOUT = 60

	// MESSAGES_QUEUE_SIZE - size of channel for bot messages
	MESSAGES_QUEUE_SIZE = 10

	// MAX_MESSAGE_LENGTH - max length of one bot message
	MAX_MESSAGE_LENGTH = 4096

	// SECONDS_FOR_AUTO_SAVE_USERS_TO_DB - save users to file every 1 min (if need)
	SECONDS_FOR_AUTO_SAVE_USERS_TO_DB = 60

	// DB_FILE_NAME - DB json name
	DB_FILE_NAME = "shell2telegram.json"
)

// Command - one user command
type Command struct {
	shellCmd    string   // shell command
	description string   // command description for list in /help (/cmd:desc="Command name")
	vars        []string // environment vars for user text, split by `/s+` to vars (/cmd:vars=SUBCOMMAND,ARGS)
}

// Commands - list of all commands
type Commands map[string]Command

// Config - config struct
type Config struct {
	token                  string   // bot token
	addExit                bool     // adding /shell2telegram exit command
	botTimeout             int      // bot timeout
	predefinedAllowedUsers []string // telegram users who are allowed to chat with the bot
	predefinedRootUsers    []string // telegram users, who confirms new users in their private chat
	allowAll               bool     // allow all user (DANGEROUS!)
	logCommands            bool     // logging all commands
	description            string   // description of bot
	persistentUsers        bool     // load/save users from file
	usersDB                string   // file for store users
}

// message types
const (
	msgIsText int8 = iota
	msgIsPhoto
)

// BotMessage - record for send via channel for send message to telegram chat
type BotMessage struct {
	chatID      int
	messageType int8
	message     string
	fileName    string
	photo       []byte
}

// ----------------------------------------------------------------------------
// get config
func getConfig() (commands Commands, appConfig Config, err error) {
	flag.StringVar(&appConfig.token, "tb-token", "", "setting bot token (or set TB_TOKEN variable)")
	flag.BoolVar(&appConfig.addExit, "add-exit", false, "adding \"/shell2telegram exit\" command for terminate bot (for roots only)")
	flag.IntVar(&appConfig.botTimeout, "timeout", DEFAULT_BOT_TIMEOUT, "setting timeout for bot")
	flag.BoolVar(&appConfig.allowAll, "allow-all", false, "allow all users (DANGEROUS!)")
	flag.BoolVar(&appConfig.logCommands, "log-commands", false, "logging all commands")
	flag.StringVar(&appConfig.description, "description", "", "setting description of bot")
	flag.BoolVar(&appConfig.persistentUsers, "persistent_users", false, "load/save users from file (default ~/.config/shell2telegram.json)")
	flag.StringVar(&appConfig.usersDB, "users_db", "", "file for store users")
	logFilename := flag.String("log", "", "log filename, default - STDOUT")
	predefinedAllowedUsers := flag.String("allow-users", "", "telegram users who are allowed to chat with the bot (\"user1,user2\")")
	predefinedRootUsers := flag.String("root-users", "", "telegram users, who confirms new users in their private chat (\"user1,user2\")")
	version := flag.Bool("version", false, "get version")

	flag.Usage = func() {
		fmt.Printf("usage: %s [options] %s\n%s\n%s\n\noptions:\n",
			os.Args[0],
			`/chat_command "shell command" /chat_command2 "shell command2"`,
			"All text after /chat_command will be sent to STDIN of shell command.",
			"If chat command is /:plain_text - get user message without any /command (for private chats only)",
		)
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.Parse()

	if *version {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	// setup log file
	if len(*logFilename) > 0 {
		fhLog, err := os.OpenFile(*logFilename, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
		if err != nil {
			log.Fatalf("error opening log file: %v", err)
		}
		log.SetOutput(fhLog)
	}

	// setup users and roots
	if *predefinedAllowedUsers != "" {
		appConfig.predefinedAllowedUsers = strings.Split(*predefinedAllowedUsers, ",")
	}
	if *predefinedRootUsers != "" {
		appConfig.predefinedRootUsers = strings.Split(*predefinedRootUsers, ",")
	}

	commands = Commands{}
	// need >= 2 arguments and count of it must be even
	args := flag.Args()
	if len(args) < 2 || len(args)%2 == 1 {
		return commands, appConfig, fmt.Errorf("error: need pairs of /chat-command and shell-command")
	}

	for i := 0; i < len(args); i += 2 {
		path, command, err := parseBotCommand(args[i], args[i+1]) // (/path, shell_command)
		if err != nil {
			return commands, appConfig, err
		}
		commands[path] = command
	}

	if appConfig.token == "" {
		if appConfig.token = os.Getenv("TB_TOKEN"); appConfig.token == "" {
			return commands, appConfig, fmt.Errorf("TB_TOKEN environment var not found. See https://core.telegram.org/bots#botfather for more information\n")
		}
	}

	return commands, appConfig, nil
}

// ----------------------------------------------------------------------------
func sendMessage(messageSignal chan<- BotMessage, chatID int, message []byte) {
	go func() {
		fileName := ""
		fileType := http.DetectContentType(message)
		switch fileType {
		case "image/png":
			fileName = "file.png"
		case "image/jpeg":
			fileName = "file.jpeg"
		case "image/gif":
			fileName = "file.gif"
		case "image/bmp":
			fileName = "file.bmp"
		default:
			fileName = "message"
		}

		if fileName == "message" {

			// is text message
			messageString := string(message)
			messagesList := []string{}

			if len(messageString) <= MAX_MESSAGE_LENGTH {
				messagesList = []string{messageString}
			} else {
				messagesList = splitStringLinesBySize(messageString, MAX_MESSAGE_LENGTH)
			}

			for _, messageChunk := range messagesList {
				messageSignal <- BotMessage{
					chatID:      chatID,
					messageType: msgIsText,
					message:     messageChunk,
				}
			}

		} else {
			// is image
			messageSignal <- BotMessage{
				chatID:      chatID,
				messageType: msgIsPhoto,
				fileName:    fileName,
				photo:       message,
			}
		}
	}()
}

// ----------------------------------------------------------------------------
func main() {
	commands, appConfig, err := getConfig()
	if err != nil {
		log.Fatal(err)
	}

	bot, err := tgbotapi.NewBotAPI(appConfig.token)
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Authorized on bot account: @%s", bot.Self.UserName)

	tgbotConfig := tgbotapi.NewUpdate(0)
	tgbotConfig.Timeout = appConfig.botTimeout
	err = bot.UpdatesChan(tgbotConfig)
	if err != nil {
		log.Fatal(err)
	}

	users := NewUsers(appConfig)
	messageSignal := make(chan BotMessage, MESSAGES_QUEUE_SIZE)
	vacuumTicker := time.Tick(SECONDS_FOR_OLD_USERS_BEFORE_VACUUM * time.Second)
	saveToBDTicker := make(<-chan time.Time)
	systemExitSignal := make(chan os.Signal)

	if appConfig.persistentUsers {
		saveToBDTicker = time.Tick(SECONDS_FOR_AUTO_SAVE_USERS_TO_DB * time.Second)
		signal.Notify(systemExitSignal, os.Interrupt, os.Kill)
	}
	exitSignal := make(chan struct{})

	// all /shell2telegram sub-commands handlers
	internalCommands := map[string]func(Ctx) string{
		"stat":              cmdShell2telegramStat,
		"ban":               cmdShell2telegramBan,
		"search":            cmdShell2telegramSearch,
		"desc":              cmdShell2telegramDesc,
		"rm":                cmdShell2telegramRm,
		"exit":              cmdShell2telegramExit,
		"version":           cmdShell2telegramVersion,
		"broadcast_to_root": cmdShell2telegramBroadcastToRoot,
		"message_to_user":   cmdShell2telegramMessageToUser,
	}

	doExit := false
	for !doExit {
		select {
		case telegramUpdate := <-bot.Updates:

			var messageCmd, messageArgs string
			allUserMessage := telegramUpdate.Message.Text
			if len(allUserMessage) > 0 && allUserMessage[0] == '/' {
				messageCmd, messageArgs = splitStringHalfBySpace(allUserMessage)
			} else {
				messageCmd, messageArgs = "/:plain_text", allUserMessage
			}

			allowPlainText := false
			if _, ok := commands["/:plain_text"]; ok {
				allowPlainText = true
			}

			replayMsg := ""

			if len(messageCmd) > 0 && (messageCmd != "/:plain_text" || allowPlainText) {

				users.AddNew(telegramUpdate.Message)
				userID := telegramUpdate.Message.From.ID
				allowExec := appConfig.allowAll || users.IsAuthorized(userID)

				ctx := Ctx{
					appConfig:     &appConfig,
					users:         &users,
					commands:      commands,
					userID:        userID,
					allowExec:     allowExec,
					messageCmd:    messageCmd,
					messageArgs:   messageArgs,
					messageSignal: messageSignal,
					chatID:        telegramUpdate.Message.Chat.ID,
					exitSignal:    exitSignal,
				}

				switch {
				// commands .................................
				case messageCmd == "/auth" || messageCmd == "/authroot":
					replayMsg = cmdAuth(ctx)

				case messageCmd == "/help":
					replayMsg = cmdHelp(ctx)

				case messageCmd == "/shell2telegram" && users.IsRoot(userID):
					messageSubCmd, messageArgs := splitStringHalfBySpace(messageArgs)
					ctx.messageArgs = messageArgs
					if cmdHandler, ok := internalCommands[messageSubCmd]; ok {
						replayMsg = cmdHandler(ctx)
					} else {
						replayMsg = "Sub-command not found"
					}

				case allowExec && (allowPlainText && messageCmd == "/:plain_text" || messageCmd[0] == '/'):
					_ = cmdUser(ctx)

				} // switch for commands

				if appConfig.logCommands {
					log.Printf("%s: %s", users.String(userID), allUserMessage)
				}

				sendMessage(messageSignal, telegramUpdate.Message.Chat.ID, []byte(replayMsg))
			}

		case botMessage := <-messageSignal:
			switch {
			case botMessage.messageType == msgIsText && !stringIsEmpty(botMessage.message):
				_, err = bot.SendMessage(tgbotapi.NewMessage(botMessage.chatID, botMessage.message))
			case botMessage.messageType == msgIsPhoto && len(botMessage.photo) > 0:
				bytesPhoto := tgbotapi.FileBytes{Name: botMessage.fileName, Bytes: botMessage.photo}
				_, err = bot.SendPhoto(tgbotapi.NewPhotoUpload(botMessage.chatID, bytesPhoto))
			}

			if err != nil {
				log.Print("Bot send message error: ", err)
			}

		case <-saveToBDTicker:
			users.SaveToDB(appConfig.usersDB)

		case <-vacuumTicker:
			users.ClearOldUsers()

		case <-systemExitSignal:
			go func() {
				exitSignal <- struct{}{}
			}()

		case <-exitSignal:
			if appConfig.persistentUsers {
				users.needSaveDB = true
				users.SaveToDB(appConfig.usersDB)
			}
			doExit = true
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"strings"
	"testing"
)

func testcmd(cmd string) (string, error) {
	switch {
	case cmd == "sudo ls -a1F /mnt/sda1/var/lib/docker/vfs/dir":
		return currenttest.vs.ls(), nil
	case cmd == "docker ps -aq --no-trunc":
		return "", nil
	case strings.HasPrefix(cmd, "docker inspect -f '{{ .Name }},{{ range $key, $value := .Volumes }}{{ $key }},{{ $value }}##~#{{ end }}' "):
		return "", nil
	case strings.HasPrefix(cmd, "sudo rm /mnt/sda1/var/lib/docker/vfs/dir/"):
		deleted := cmd[len("sudo rm /mnt/sda1/var/lib/docker/vfs/dir/"):]
		deletions = append(deletions, deleted)
		return "", nil
	case strings.HasPrefix(cmd, "sudo readlink /mnt/sda1/var/lib/docker/vfs/dir/"):
		if strings.Contains(cmd, ",nonexistent") {
			return "", errors.New("non-existent linked folder")
		}
		return "", nil
	case strings.HasPrefix(cmd, "sudo ls /mnt/sda1/var/lib/docker/vfs/dir/"):
		if cmd == "sudo ls /mnt/sda1/var/lib/docker/vfs/dir/" {
			return "", errors.New("non-ls linked folder")
		}
		return "", nil
	default:
		currentT.Fatalf("test '%s': unknown command!\n", cmd)
		return fmt.Sprintf("test '%s'", cmd), errors.New("unknown command")
	}
}

type volspecs []string
type Test struct {
	title string
	vs    volspecs
	res   []int
}

func (vs volspecs) ls() string {
	if len(vs) == 0 {
		return ""
	}
	res := ""
	for i, spec := range vs {
		if strings.HasSuffix(spec, "/") {
			spec = spec[:len(spec)-1]
			res = res + spec + strings.Repeat(fmt.Sprintf("%d", i), 64-len(spec)) + "/\n"
		}
		if strings.HasSuffix(spec, "@") {
			mp := "." + strings.Replace(spec, ";", "###", -1)
			mp = strings.Replace(mp, "/", ",#,", -1)
			res = res + mp + "\n"
		}
	}
	return res
}

var deletions = []string{}
var tests = []Test{
	Test{"empty vfs", []string{}, []int{0, 0, 0, 0, 0}},
	Test{"two volumes", []string{"fa/", "fb/"}, []int{0, 0, 2, 2, 0}},
	Test{"Invalid (ill-formed) markers must be deleted", []string{"cainv/path/a@"}, []int{0, 0, 0, 0, -1}},
	Test{"Invalid (no readlink) markers must be deleted", []string{"ca;/path/nonexistenta@", "cb;/path/nonexistentb@"}, []int{0, 0, 0, 0, -2}},
	Test{"Invalid (no ls) markers must be deleted", []string{"ca;/path/nolsa@", "cb;/path/nolsb@"}, []int{0, 0, 0, 0, -2}},
}
var currenttest Test
var currentT *testing.T

// TestContainers test different vfs scenarios
func TestContainers(t *testing.T) {
	cmd = testcmd
	currentT = t
	for i, test := range tests {
		currenttest = test
		deletions = []string{}
		main()
		tc := Containers()
		toc := OrphanedContainers()
		tv := Volumes()
		tov := OrphanedVolumes()
		tm := Markers()
		if len(tc) != test.res[0] {
			t.Errorf("Test %d: '%s' expected '%d' containers, got '%d'", i+1, test.title, test.res[0], len(tc))
		}
		if len(toc) != test.res[1] {
			t.Errorf("Test %d: '%s' expected '%d' orphaned containers, got '%d'", i+1, test.title, test.res[1], len(toc))
		}
		if len(tv) != test.res[2] {
			t.Errorf("Test %d: '%s' expected '%d' volumes, got '%d'", i+1, test.title, test.res[2], len(tv))
		}
		if len(tov) != test.res[3] {
			t.Errorf("Test %d: '%s' expected '%d' orphaned volumes, got '%d'", i+1, test.title, test.res[3], len(tov))
		}
		if nbmarkers(tm) != test.res[4] {
			t.Errorf("Test %d: '%s' expected '%d' markers, got '%d'", i+1, test.title, test.res[4], nbmarkers(tm))
		}
		fmt.Println("----------------")
	}
}

func nbmarkers(tm markers) int {
	res := len(tm)
	for _, d := range deletions {
		if strings.HasPrefix(d, ".") {
			res = res - 1
		}
	}
	return res
}
<commit_msg>gcl_test.go: display testcmd cmds<commit_after>package main

import (
	"errors"
	"fmt"
	"strings"
	"testing"
)

func testcmd(cmd string) (string, error) {
	fmt.Println("(T) " + cmd)
	switch {
	case cmd == "sudo ls -a1F /mnt/sda1/var/lib/docker/vfs/dir":
		return currenttest.vs.ls(), nil
	case cmd == "docker ps -aq --no-trunc":
		return "", nil
	case strings.HasPrefix(cmd, "docker inspect -f '{{ .Name }},{{ range $key, $value := .Volumes }}{{ $key }},{{ $value }}##~#{{ end }}' "):
		return "", nil
	case strings.HasPrefix(cmd, "sudo rm /mnt/sda1/var/lib/docker/vfs/dir/"):
		deleted := cmd[len("sudo rm /mnt/sda1/var/lib/docker/vfs/dir/"):]
		deletions = append(deletions, deleted)
		return "", nil
	case strings.HasPrefix(cmd, "sudo readlink /mnt/sda1/var/lib/docker/vfs/dir/"):
		if strings.Contains(cmd, ",nonexistent") {
			return "", errors.New("non-existent linked folder")
		}
		return "", nil
	case strings.HasPrefix(cmd, "sudo ls /mnt/sda1/var/lib/docker/vfs/dir/"):
		if cmd == "sudo ls /mnt/sda1/var/lib/docker/vfs/dir/" {
			return "", errors.New("non-ls linked folder")
		}
		return "", nil
	default:
		currentT.Fatalf("test '%s': unknown command!\n", cmd)
		return fmt.Sprintf("test '%s'", cmd), errors.New("unknown command")
	}
}

type volspecs []string
type Test struct {
	title string
	vs    volspecs
	res   []int
}

func (vs volspecs) ls() string {
	if len(vs) == 0 {
		return ""
	}
	res := ""
	for i, spec := range vs {
		if strings.HasSuffix(spec, "/") {
			spec = spec[:len(spec)-1]
			res = res + spec + strings.Repeat(fmt.Sprintf("%d", i), 64-len(spec)) + "/\n"
		}
		if strings.HasSuffix(spec, "@") {
			mp := "." + strings.Replace(spec, ";", "###", -1)
			mp = strings.Replace(mp, "/", ",#,", -1)
			res = res + mp + "\n"
		}
	}
	return res
}

var deletions = []string{}
var tests = []Test{
	Test{"empty vfs", []string{}, []int{0, 0, 0, 0, 0}},
	Test{"two volumes", []string{"fa/", "fb/"}, []int{0, 0, 2, 2, 0}},
	Test{"Invalid (ill-formed) markers must be deleted", []string{"cainv/path/a@"}, []int{0, 0, 0, 0, -1}},
	Test{"Invalid (no readlink) markers must be deleted", []string{"ca;/path/nonexistenta@", "cb;/path/nonexistentb@"}, []int{0, 0, 0, 0, -2}},
	Test{"Invalid (no ls) markers must be deleted", []string{"ca;/path/nolsa@", "cb;/path/nolsb@"}, []int{0, 0, 0, 0, -2}},
}
var currenttest Test
var currentT *testing.T

// TestContainers test different vfs scenarios
func TestContainers(t *testing.T) {
	cmd = testcmd
	currentT = t
	for i, test := range tests {
		currenttest = test
		deletions = []string{}
		main()
		tc := Containers()
		toc := OrphanedContainers()
		tv := Volumes()
		tov := OrphanedVolumes()
		tm := Markers()
		if len(tc) != test.res[0] {
			t.Errorf("Test %d: '%s' expected '%d' containers, got '%d'", i+1, test.title, test.res[0], len(tc))
		}
		if len(toc) != test.res[1] {
			t.Errorf("Test %d: '%s' expected '%d' orphaned containers, got '%d'", i+1, test.title, test.res[1], len(toc))
		}
		if len(tv) != test.res[2] {
			t.Errorf("Test %d: '%s' expected '%d' volumes, got '%d'", i+1, test.title, test.res[2], len(tv))
		}
		if len(tov) != test.res[3] {
			t.Errorf("Test %d: '%s' expected '%d' orphaned volumes, got '%d'", i+1, test.title, test.res[3], len(tov))
		}
		if nbmarkers(tm) != test.res[4] {
			t.Errorf("Test %d: '%s' expected '%d' markers, got '%d'", i+1, test.title, test.res[4], nbmarkers(tm))
		}
		fmt.Println("----------------")
	}
}

func nbmarkers(tm markers) int {
	res := len(tm)
	for _, d := range deletions {
		if strings.HasPrefix(d, ".") {
			res = res - 1
		}
	}
	return res
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"errors"
	"crypto/rand"
	"math/big"
)

// generate creates a random sequence of valid characters of a certain length
func generate(length int, valid charset) (string, error) {
	pass := ""
	errs := []error{}
	for i := 0; i < length; i++ {
		// Attempt to get a random element
		char, err := randElem(valid)
		if err != nil {
			// Record error
			errs := append(errs, err)
			// Tolerate up to 50% error rate, and at least 5
			if len(errs) > max(length / 2, 5) {
				return "", tooMany(errs)
			}
		} else {
			// If there is no error, add the rune to the return string
			pass += fmt.Sprintf("%c", char)
		}
	}
	return pass, nil
}

// randElem gets a random rune from a charset
func randElem(set charset) (rune, error) {
	// Create a list to choose a random index from
	list := getlist(set)

	// Set the maximum index to choose - casting to big int for crypto/rand
	max := big.NewInt(int64(len(list)))

	// Generate a random index (See godoc for crypto/rand for info)
	i, err := rand.Int(rand.Reader, max)

	// Error reading from os random source
	if err != nil {
		return ' ', err
	}

	// No error, return rune at random index - casting back to int
	return list[int(i.Int64())], nil
}

// max gets the maximum value of a list of values
func max(is ...int) int {
	m := 0
	for i := range is {
		if i > m {
			m = i
		}
	}
	return m
}

// tooMany returns an amalgamated error message listing all of the
// errors that caused the calling function to give up
func tooMany(errs []error) error {
	message := "Too many errors: \n"
	for i, e := range errs {
		message += fmt.Sprintln("\t", i, "-", e)
	}
	return errors.New(message)
}
<commit_msg>Adds a comment explaining the odd behavior of Int<commit_after>package main

import (
	"fmt"
	"errors"
	"crypto/rand"
	"math/big"
)

// generate creates a random sequence of valid characters of a certain length
func generate(length int, valid charset) (string, error) {
	pass := ""
	errs := []error{}
	for i := 0; i < length; i++ {
		// Attempt to get a random element
		char, err := randElem(valid)
		if err != nil {
			// Record error
			errs := append(errs, err)
			// Tolerate up to 50% error rate, and at least 5
			if len(errs) > max(length / 2, 5) {
				return "", tooMany(errs)
			}
		} else {
			// If there is no error, add the rune to the return string
			pass += fmt.Sprintf("%c", char)
		}
	}
	return pass, nil
}

// randElem gets a random rune from a charset
func randElem(set charset) (rune, error) {
	// Create a list to choose a random index from
	list := getlist(set)

	// Set the maximum index to choose - casting to big int for crypto/rand
	max := big.NewInt(int64(len(list)))

	// Generate a random index
	// crypto/rand.Int generates a number in [0,max) - See godoc for crypto/rand for more info
	i, err := rand.Int(rand.Reader, max)

	// Error reading from os random source
	if err != nil {
		return ' ', err
	}

	// No error, return rune at random index - casting back to int
	return list[int(i.Int64())], nil
}

// max gets the maximum value of a list of values
func max(is ...int) int {
	m := 0
	for i := range is {
		if i > m {
			m = i
		}
	}
	return m
}

// tooMany returns an amalgamated error message listing all of the
// errors that caused the calling function to give up
func tooMany(errs []error) error {
	message := "Too many errors: \n"
	for i, e := range errs {
		message += fmt.Sprintln("\t", i, "-", e)
	}
	return errors.New(message)
}
<|endoftext|>


<commit_before>package azuread

import (
	"fmt"
	"log"

	"github.com/Azure/azure-sdk-for-go/services/graphrbac/1.6/graphrbac"
	"github.com/google/uuid"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/helper/validation"

	"github.com/terraform-providers/terraform-provider-azuread/azuread/helpers/ar"
	"github.com/terraform-providers/terraform-provider-azuread/azuread/helpers/graph"
	"github.com/terraform-providers/terraform-provider-azuread/azuread/helpers/p"
	"github.com/terraform-providers/terraform-provider-azuread/azuread/helpers/slices"
	"github.com/terraform-providers/terraform-provider-azuread/azuread/helpers/tf"
	"github.com/terraform-providers/terraform-provider-azuread/azuread/helpers/validate"
)

func resourceGroup() *schema.Resource {
	return &schema.Resource{
		Create: resourceGroupCreate,
		Read:   resourceGroupRead,
		Update: resourceGroupUpdate,
		Delete: resourceGroupDelete,

		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"name": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validation.NoZeroValues,
			},

			"description": {
				Type:     schema.TypeString,
				ForceNew: true, // there is no update method available in the SDK
				Optional: true,
			},

			"members": {
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Set:      schema.HashString,
				Elem: &schema.Schema{
					Type:         schema.TypeString,
					ValidateFunc: validate.UUID,
				},
			},

			"owners": {
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Set:      schema.HashString,
				Elem: &schema.Schema{
					Type:         schema.TypeString,
					ValidateFunc: validate.UUID,
				},
			},

			"object_id": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceGroupCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient).groupsClient
	ctx := meta.(*ArmClient).StopContext

	name := d.Get("name").(string)

	properties := graphrbac.GroupCreateParameters{
		DisplayName:          &name,
		MailEnabled:          p.Bool(false),                 // we're defaulting to false, as the API currently only supports the creation of non-mail enabled security groups.
		MailNickname:         p.String(uuid.New().String()), // this matches the portal behaviour
		SecurityEnabled:      p.Bool(true),                  // we're defaulting to true, as the API currently only supports the creation of non-mail enabled security groups.
		AdditionalProperties: make(map[string]interface{}),
	}

	if v, ok := d.GetOk("description"); ok {
		properties.AdditionalProperties["description"] = v.(string)
	}

	group, err := client.Create(ctx, properties)
	if err != nil {
		return fmt.Errorf("Error creating Group (%q): %+v", name, err)
	}
	if group.ObjectID == nil {
		return fmt.Errorf("nil Group ID for %q: %+v", name, err)
	}

	d.SetId(*group.ObjectID)

	// Add members if specified
	if v, ok := d.GetOk("members"); ok {
		members := tf.ExpandStringSlicePtr(v.(*schema.Set).List())

		// we could lock here against the group member resource, but they should not be used together (todo conflicts with at a resource level?)
		if err := graph.GroupAddMembers(client, ctx, *group.ObjectID, *members); err != nil {
			return err
		}
	}

	// Add owners if specified
	if v, ok := d.GetOk("owners"); ok {
		members := tf.ExpandStringSlicePtr(v.(*schema.Set).List())
		if err := graph.GroupAddOwners(client, ctx, *group.ObjectID, *members); err != nil {
			return err
		}
	}

	_, err = graph.WaitForCreationReplication(func() (interface{}, error) {
		return client.Get(ctx, *group.ObjectID)
	})
	if err != nil {
		return fmt.Errorf("Error waiting for Group (%s) with ObjectId %q: %+v", name, *group.ObjectID, err)
	}

	return resourceGroupRead(d, meta)
}

func resourceGroupRead(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient).groupsClient
	ctx := meta.(*ArmClient).StopContext

	resp, err := client.Get(ctx, d.Id())
	if err != nil {
		if ar.ResponseWasNotFound(resp.Response) {
			log.Printf("[DEBUG] Azure AD group with id %q was not found - removing from state", d.Id())
			d.SetId("")
			return nil
		}

		return fmt.Errorf("Error retrieving Azure AD Group with ID %q: %+v", d.Id(), err)
	}

	d.Set("name", resp.DisplayName)
	d.Set("object_id", resp.ObjectID)

	if v, ok := resp.AdditionalProperties["description"]; ok {
		d.Set("description", v.(string))
	}

	members, err := graph.GroupAllMembers(client, ctx, d.Id())
	if err != nil {
		return err
	}
	d.Set("members", members)

	owners, err := graph.GroupAllOwners(client, ctx, d.Id())
	if err != nil {
		return err
	}
	d.Set("owners", owners)

	return nil
}

func resourceGroupUpdate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient).groupsClient
	ctx := meta.(*ArmClient).StopContext

	if v, ok := d.GetOkExists("members"); ok && d.HasChange("members") {
		existingMembers, err := graph.GroupAllMembers(client, ctx, d.Id())
		if err != nil {
			return err
		}

		desiredMembers := *tf.ExpandStringSlicePtr(v.(*schema.Set).List())
		membersForRemoval := slices.Difference(existingMembers, desiredMembers)
		membersToAdd := slices.Difference(desiredMembers, existingMembers)

		for _, existingMember := range membersForRemoval {
			log.Printf("[DEBUG] Removing member with id %q from Azure AD group with id %q", existingMember, d.Id())
			if resp, err := client.RemoveMember(ctx, d.Id(), existingMember); err != nil {
				if !ar.ResponseWasNotFound(resp) {
					return fmt.Errorf("Error Deleting group member %q from Azure AD Group with ID %q: %+v", existingMember, d.Id(), err)
				}
			}
		}

		if err := graph.GroupAddMembers(client, ctx, d.Id(), membersToAdd); err != nil {
			return err
		}
	}

	if v, ok := d.GetOkExists("owners"); ok && d.HasChange("owners") {
		existingOwners, err := graph.GroupAllOwners(client, ctx, d.Id())
		if err != nil {
			return err
		}

		desiredOwners := *tf.ExpandStringSlicePtr(v.(*schema.Set).List())
		ownersForRemoval := slices.Difference(existingOwners, desiredOwners)
		ownersToAdd := slices.Difference(desiredOwners, existingOwners)

		for _, ownerToDelete := range ownersForRemoval {
			log.Printf("[DEBUG] Removing member with id %q from Azure AD group with id %q", ownerToDelete, d.Id())
			if resp, err := client.RemoveOwner(ctx, d.Id(), ownerToDelete); err != nil {
				if !ar.ResponseWasNotFound(resp) {
					return fmt.Errorf("Error Deleting group member %q from Azure AD Group with ID %q: %+v", ownerToDelete, d.Id(), err)
				}
			}
		}

		if err := graph.GroupAddOwners(client, ctx, d.Id(), ownersToAdd); err != nil {
			return err
		}
	}

	return resourceGroupRead(d, meta)
}

func resourceGroupDelete(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient).groupsClient
	ctx := meta.(*ArmClient).StopContext

	if resp, err := client.Delete(ctx, d.Id()); err != nil {
		if !ar.ResponseWasNotFound(resp) {
			return fmt.Errorf("Error Deleting Azure AD Group with ID %q: %+v", d.Id(), err)
		}
	}

	return nil
}
<commit_msg>azuread_group - owners is now additive on creation allowing exis… (#211)<commit_after>package azuread

import (
	"fmt"
	"log"

	"github.com/Azure/azure-sdk-for-go/services/graphrbac/1.6/graphrbac"
	"github.com/google/uuid"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/helper/validation"

	"github.com/terraform-providers/terraform-provider-azuread/azuread/helpers/ar"
	"github.com/terraform-providers/terraform-provider-azuread/azuread/helpers/graph"
	"github.com/terraform-providers/terraform-provider-azuread/azuread/helpers/p"
	"github.com/terraform-providers/terraform-provider-azuread/azuread/helpers/slices"
	"github.com/terraform-providers/terraform-provider-azuread/azuread/helpers/tf"
	"github.com/terraform-providers/terraform-provider-azuread/azuread/helpers/validate"
)

func resourceGroup() *schema.Resource {
	return &schema.Resource{
		Create: resourceGroupCreate,
		Read:   resourceGroupRead,
		Update: resourceGroupUpdate,
		Delete: resourceGroupDelete,

		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"name": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validation.NoZeroValues,
			},

			"description": {
				Type:     schema.TypeString,
				ForceNew: true, // there is no update method available in the SDK
				Optional: true,
			},

			"members": {
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Set:      schema.HashString,
				Elem: &schema.Schema{
					Type:         schema.TypeString,
					ValidateFunc: validate.UUID,
				},
			},

			"owners": {
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Set:      schema.HashString,
				Elem: &schema.Schema{
					Type:         schema.TypeString,
					ValidateFunc: validate.UUID,
				},
			},

			"object_id": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceGroupCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient).groupsClient
	ctx := meta.(*ArmClient).StopContext

	name := d.Get("name").(string)

	properties := graphrbac.GroupCreateParameters{
		DisplayName:          &name,
		MailEnabled:          p.Bool(false),                 // we're defaulting to false, as the API currently only supports the creation of non-mail enabled security groups.
		MailNickname:         p.String(uuid.New().String()), // this matches the portal behaviour
		SecurityEnabled:      p.Bool(true),                  // we're defaulting to true, as the API currently only supports the creation of non-mail enabled security groups.
		AdditionalProperties: make(map[string]interface{}),
	}

	if v, ok := d.GetOk("description"); ok {
		properties.AdditionalProperties["description"] = v.(string)
	}

	group, err := client.Create(ctx, properties)
	if err != nil {
		return fmt.Errorf("Error creating Group (%q): %+v", name, err)
	}
	if group.ObjectID == nil {
		return fmt.Errorf("nil Group ID for %q: %+v", name, err)
	}

	d.SetId(*group.ObjectID)

	// Add members if specified
	if v, ok := d.GetOk("members"); ok {
		members := tf.ExpandStringSlicePtr(v.(*schema.Set).List())

		// we could lock here against the group member resource, but they should not be used together (todo conflicts with at a resource level?)
		if err := graph.GroupAddMembers(client, ctx, *group.ObjectID, *members); err != nil {
			return err
		}
	}

	// Add owners if specified
	if v, ok := d.GetOk("owners"); ok {
		existingOwners, err := graph.GroupAllOwners(client, ctx, *group.ObjectID)
		if err != nil {
			return err
		}
		members := *tf.ExpandStringSlicePtr(v.(*schema.Set).List())
		ownersToAdd := slices.Difference(members, existingOwners)

		if err := graph.GroupAddOwners(client, ctx, *group.ObjectID, ownersToAdd); err != nil {
			return err
		}
	}

	_, err = graph.WaitForCreationReplication(func() (interface{}, error) {
		return client.Get(ctx, *group.ObjectID)
	})
	if err != nil {
		return fmt.Errorf("Error waiting for Group (%s) with ObjectId %q: %+v", name, *group.ObjectID, err)
	}

	return resourceGroupRead(d, meta)
}

func resourceGroupRead(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient).groupsClient
	ctx := meta.(*ArmClient).StopContext

	resp, err := client.Get(ctx, d.Id())
	if err != nil {
		if ar.ResponseWasNotFound(resp.Response) {
			log.Printf("[DEBUG] Azure AD group with id %q was not found - removing from state", d.Id())
			d.SetId("")
			return nil
		}

		return fmt.Errorf("Error retrieving Azure AD Group with ID %q: %+v", d.Id(), err)
	}

	d.Set("name", resp.DisplayName)
	d.Set("object_id", resp.ObjectID)

	if v, ok := resp.AdditionalProperties["description"]; ok {
		d.Set("description", v.(string))
	}

	members, err := graph.GroupAllMembers(client, ctx, d.Id())
	if err != nil {
		return err
	}
	d.Set("members", members)

	owners, err := graph.GroupAllOwners(client, ctx, d.Id())
	if err != nil {
		return err
	}
	d.Set("owners", owners)

	return nil
}

func resourceGroupUpdate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient).groupsClient
	ctx := meta.(*ArmClient).StopContext

	if v, ok := d.GetOkExists("members"); ok && d.HasChange("members") {
		existingMembers, err := graph.GroupAllMembers(client, ctx, d.Id())
		if err != nil {
			return err
		}

		desiredMembers := *tf.ExpandStringSlicePtr(v.(*schema.Set).List())
		membersForRemoval := slices.Difference(existingMembers, desiredMembers)
		membersToAdd := slices.Difference(desiredMembers, existingMembers)

		for _, existingMember := range membersForRemoval {
			log.Printf("[DEBUG] Removing member with id %q from Azure AD group with id %q", existingMember, d.Id())
			if resp, err := client.RemoveMember(ctx, d.Id(), existingMember); err != nil {
				if !ar.ResponseWasNotFound(resp) {
					return fmt.Errorf("Error Deleting group member %q from Azure AD Group with ID %q: %+v", existingMember, d.Id(), err)
				}
			}
		}

		if err := graph.GroupAddMembers(client, ctx, d.Id(), membersToAdd); err != nil {
			return err
		}
	}

	if v, ok := d.GetOkExists("owners"); ok && d.HasChange("owners") {
		existingOwners, err := graph.GroupAllOwners(client, ctx, d.Id())
		if err != nil {
			return err
		}

		desiredOwners := *tf.ExpandStringSlicePtr(v.(*schema.Set).List())
		ownersForRemoval := slices.Difference(existingOwners, desiredOwners)
		ownersToAdd := slices.Difference(desiredOwners, existingOwners)

		for _, ownerToDelete := range ownersForRemoval {
			log.Printf("[DEBUG] Removing member with id %q from Azure AD group with id %q", ownerToDelete, d.Id())
			if resp, err := client.RemoveOwner(ctx, d.Id(), ownerToDelete); err != nil {
				if !ar.ResponseWasNotFound(resp) {
					return fmt.Errorf("Error Deleting group member %q from Azure AD Group with ID %q: %+v", ownerToDelete, d.Id(), err)
				}
			}
		}

		if err := graph.GroupAddOwners(client, ctx, d.Id(), ownersToAdd); err != nil {
			return err
		}
	}

	return resourceGroupRead(d, meta)
}

func resourceGroupDelete(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*ArmClient).groupsClient
	ctx := meta.(*ArmClient).StopContext

	if resp, err := client.Delete(ctx, d.Id()); err != nil {
		if !ar.ResponseWasNotFound(resp) {
			return fmt.Errorf("Error Deleting Azure AD Group with ID %q: %+v", d.Id(), err)
		}
	}

	return nil
}
<|endoftext|>