<commit_before>package storage

import (
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"sync"
)

const (
	SuperBlockSize = 8
)

type SuperBlock struct {
	Version     Version
	ReplicaType ReplicationType
}

func (s *SuperBlock) Bytes() []byte {
	header := make([]byte, SuperBlockSize)
	header[0] = byte(s.Version)
	header[1] = s.ReplicaType.Byte()
	return header
}

type Volume struct {
	Id       VolumeId
	dir      string
	dataFile *os.File
	nm       NeedleMapper
	readOnly bool

	SuperBlock

	accessLock sync.Mutex
}

func NewVolume(dirname string, id VolumeId, replicationType ReplicationType) (v *Volume, e error) {
	v = &Volume{dir: dirname, Id: id}
	v.SuperBlock = SuperBlock{ReplicaType: replicationType}
	e = v.load(true)
	return
}
func LoadVolumeOnly(dirname string, id VolumeId) (v *Volume, e error) {
	v = &Volume{dir: dirname, Id: id}
	v.SuperBlock = SuperBlock{ReplicaType: CopyNil}
	e = v.load(false)
	return
}
func (v *Volume) load(alsoLoadIndex bool) error {
	var e error
	fileName := path.Join(v.dir, v.Id.String())
	v.dataFile, e = os.OpenFile(fileName+".dat", os.O_RDWR|os.O_CREATE, 0644)
	if e != nil {
		if !os.IsPermission(e) {
			return fmt.Errorf("cannot create Volume Data %s.dat: %s", fileName, e)
		}
		if v.dataFile, e = os.Open(fileName + ".dat"); e != nil {
			return fmt.Errorf("cannot open Volume Data %s.dat: %s", fileName, e)
		}
		log.Printf("opening " + fileName + ".dat in READONLY mode")
		v.readOnly = true
	}
	if v.ReplicaType == CopyNil {
		e = v.readSuperBlock()
	} else {
		e = v.maybeWriteSuperBlock()
	}
	if e == nil && alsoLoadIndex {
		var indexFile *os.File
		if v.readOnly {
			if indexFile, e = os.Open(fileName + ".idx"); e != nil && !os.IsNotExist(e) {
				return fmt.Errorf("cannot open index file %s.idx: %s", fileName, e)
			}
			if indexFile != nil {
				log.Printf("converting %s.idx to %s.cdb", fileName, fileName)
				if e = ConvertIndexToCdb(fileName+".cdb", indexFile); e != nil {
					log.Printf("error converting %s.idx to %s.cdb: %s", fileName, fileName)
				} else {
					indexFile.Close()
					os.Remove(indexFile.Name())
					indexFile = nil
				}
			}
			v.nm, e = OpenCdbMap(fileName + ".cdb")
			return e
		} else {
			indexFile, e = os.OpenFile(fileName+".idx", os.O_RDWR|os.O_CREATE, 0644)
			if e != nil {
				return fmt.Errorf("cannot create Volume Data %s.dat: %s", fileName, e)
			}
		}
		v.nm, e = LoadNeedleMap(indexFile)
	}
	return e
}
func (v *Volume) Version() Version {
	return v.SuperBlock.Version
}
func (v *Volume) Size() int64 {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	stat, e := v.dataFile.Stat()
	if e == nil {
		return stat.Size()
	}
	fmt.Printf("Failed to read file size %s %s\n", v.dataFile.Name(), e.Error())
	return -1
}
func (v *Volume) Close() {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	v.nm.Close()
	_ = v.dataFile.Close()
}
func (v *Volume) maybeWriteSuperBlock() error {
	stat, e := v.dataFile.Stat()
	if e != nil {
		fmt.Printf("failed to stat datafile %s: %s", v.dataFile, e)
		return e
	}
	if stat.Size() == 0 {
		v.SuperBlock.Version = CurrentVersion
		_, e = v.dataFile.Write(v.SuperBlock.Bytes())
		if e != nil && os.IsPermission(e) {
			//read-only, but zero length - recreate it!
			if v.dataFile, e = os.Create(v.dataFile.Name()); e == nil {
				if _, e = v.dataFile.Write(v.SuperBlock.Bytes()); e == nil {
					v.readOnly = false
				}
			}
		}
	}
	return e
}
func (v *Volume) readSuperBlock() (err error) {
	if _, err = v.dataFile.Seek(0, 0); err != nil {
		return fmt.Errorf("cannot seek to the beginning of %s: %s", v.dataFile, err)
	}
	header := make([]byte, SuperBlockSize)
	if _, e := v.dataFile.Read(header); e != nil {
		return fmt.Errorf("cannot read superblock: %s", e)
	}
	v.SuperBlock, err = ParseSuperBlock(header)
	return err
}
func ParseSuperBlock(header []byte) (superBlock SuperBlock, err error) {
	superBlock.Version = Version(header[0])
	if superBlock.ReplicaType, err = NewReplicationTypeFromByte(header[1]); err != nil {
		err = fmt.Errorf("cannot read replica type: %s", err)
	}
	return
}
func (v *Volume) NeedToReplicate() bool {
	return v.ReplicaType.GetCopyCount() > 1
}

func (v *Volume) write(n *Needle) (size uint32, err error) {
	if v.readOnly {
		err = fmt.Errorf("%s is read-only", v.dataFile)
		return
	}
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	var offset int64
	if offset, err = v.dataFile.Seek(0, 2); err != nil {
		return
	}
	if size, err = n.Append(v.dataFile, v.Version()); err != nil {
		if e := v.dataFile.Truncate(offset); e != nil {
			err = fmt.Errorf("%s\ncannot truncate %s: %s", err, v.dataFile, e)
		}
		return
	}
	nv, ok := v.nm.Get(n.Id)
	if !ok || int64(nv.Offset)*NeedlePaddingSize < offset {
		_, err = v.nm.Put(n.Id, uint32(offset/NeedlePaddingSize), n.Size)
	}
	return
}
func (v *Volume) delete(n *Needle) (uint32, error) {
	if v.readOnly {
		return 0, fmt.Errorf("%s is read-only", v.dataFile)
	}
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	nv, ok := v.nm.Get(n.Id)
	//fmt.Println("key", n.Id, "volume offset", nv.Offset, "data_size", n.Size, "cached size", nv.Size)
	if ok {
		var err error
		if err = v.nm.Delete(n.Id); err != nil {
			return nv.Size, err
		}
		if _, err = v.dataFile.Seek(int64(nv.Offset*NeedlePaddingSize), 0); err != nil {
			return nv.Size, err
		}
		_, err = n.Append(v.dataFile, v.Version())
		return nv.Size, err
	}
	return 0, nil
}

func (v *Volume) read(n *Needle) (int, error) {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	nv, ok := v.nm.Get(n.Id)
	if ok && nv.Offset > 0 {
		if _, err := v.dataFile.Seek(int64(nv.Offset)*NeedlePaddingSize, 0); err != nil {
			return -1, err
		}
		return n.Read(v.dataFile, nv.Size, v.Version())
	}
	return -1, errors.New("Not Found")
}

func (v *Volume) garbageLevel() float64 {
	return float64(v.nm.DeletedSize()) / float64(v.ContentSize())
}

func (v *Volume) compact() error {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()

	filePath := path.Join(v.dir, v.Id.String())
	return v.copyDataAndGenerateIndexFile(filePath+".cpd", filePath+".cpx")
}
func (v *Volume) commitCompact() error {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	_ = v.dataFile.Close()
	var e error
	if e = os.Rename(path.Join(v.dir, v.Id.String()+".cpd"), path.Join(v.dir, v.Id.String()+".dat")); e != nil {
		return e
	}
	if e = os.Rename(path.Join(v.dir, v.Id.String()+".cpx"), path.Join(v.dir, v.Id.String()+".idx")); e != nil {
		return e
	}
	if e = v.load(true); e != nil {
		return e
	}
	return nil
}
func (v *Volume) freeze() error {
	if v.readOnly {
		return nil
	}
	nm, ok := v.nm.(*NeedleMap)
	if !ok {
		return nil
	}
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	bn, _ := nakeFilename(v.dataFile.Name())
	cdbFn := bn + ".cdb"
	log.Printf("converting %s to %s", nm.indexFile.Name(), cdbFn)
	err := DumpNeedleMapToCdb(cdbFn, nm)
	if err != nil {
		return err
	}
	if v.nm, err = OpenCdbMap(cdbFn); err != nil {
		return err
	}
	nm.indexFile.Close()
	os.Remove(nm.indexFile.Name())
	v.readOnly = true
	return nil
}

func ScanVolumeFile(dirname string, id VolumeId,
	visitSuperBlock func(SuperBlock) error,
	visitNeedle func(n *Needle, offset uint32) error) (err error) {
	var v *Volume
	if v, err = LoadVolumeOnly(dirname, id); err != nil {
		return
	}
	if err = visitSuperBlock(v.SuperBlock); err != nil {
		return
	}

	version := v.Version()

	offset := uint32(SuperBlockSize)
	n, rest, e := ReadNeedleHeader(v.dataFile, version)
	if e != nil {
		err = fmt.Errorf("cannot read needle header: %s", e)
		return
	}
	for n != nil {
		if err = n.ReadNeedleBody(v.dataFile, version, rest); err != nil {
			err = fmt.Errorf("cannot read needle body: %s", err)
			return
		}
		if err = visitNeedle(n, offset); err != nil {
			return
		}
		offset += NeedleHeaderSize + rest
		if n, rest, err = ReadNeedleHeader(v.dataFile, version); err != nil {
			if err == io.EOF {
				return nil
			}
			return fmt.Errorf("cannot read needle header: %s", err)
		}
	}

	return
}

func (v *Volume) copyDataAndGenerateIndexFile(dstName, idxName string) (err error) {
	var (
		dst, idx *os.File
	)
	if dst, err = os.OpenFile(dstName, os.O_WRONLY|os.O_CREATE, 0644); err != nil {
		return
	}
	defer dst.Close()

	if idx, err = os.OpenFile(idxName, os.O_WRONLY|os.O_CREATE, 0644); err != nil {
		return
	}
	defer idx.Close()

	nm := NewNeedleMap(idx)
	new_offset := uint32(SuperBlockSize)

	err = ScanVolumeFile(v.dir, v.Id, func(superBlock SuperBlock) error {
		_, err = dst.Write(superBlock.Bytes())
		return err
	}, func(n *Needle, offset uint32) error {
		nv, ok := v.nm.Get(n.Id)
		//log.Println("file size is", n.Size, "rest", rest)
		if ok && nv.Offset*NeedlePaddingSize == offset {
			if nv.Size > 0 {
				if _, err = nm.Put(n.Id, new_offset/NeedlePaddingSize, n.Size); err != nil {
					return fmt.Errorf("cannot put needle: %s", err)
				}
				if _, err = n.Append(dst, v.Version()); err != nil {
					return fmt.Errorf("cannot append needle: %s", err)
				}
				new_offset += n.DiskSize()
				//log.Println("saving key", n.Id, "volume offset", old_offset, "=>", new_offset, "data_size", n.Size, "rest", rest)
			}
		}
		return nil
	})

	return
}
func (v *Volume) ContentSize() uint64 {
	return v.nm.ContentSize()
}
<commit_msg>Issue 26:	Could not download uploaded files<commit_after>package storage

import (
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"sync"
)

const (
	SuperBlockSize = 8
)

type SuperBlock struct {
	Version     Version
	ReplicaType ReplicationType
}

func (s *SuperBlock) Bytes() []byte {
	header := make([]byte, SuperBlockSize)
	header[0] = byte(s.Version)
	header[1] = s.ReplicaType.Byte()
	return header
}

type Volume struct {
	Id       VolumeId
	dir      string
	dataFile *os.File
	nm       NeedleMapper
	readOnly bool

	SuperBlock

	accessLock sync.Mutex
}

func NewVolume(dirname string, id VolumeId, replicationType ReplicationType) (v *Volume, e error) {
	v = &Volume{dir: dirname, Id: id}
	v.SuperBlock = SuperBlock{ReplicaType: replicationType}
	e = v.load(true)
	return
}
func LoadVolumeOnly(dirname string, id VolumeId) (v *Volume, e error) {
	v = &Volume{dir: dirname, Id: id}
	v.SuperBlock = SuperBlock{ReplicaType: CopyNil}
	e = v.load(false)
	return
}
func (v *Volume) load(alsoLoadIndex bool) error {
	var e error
	fileName := path.Join(v.dir, v.Id.String())
	v.dataFile, e = os.OpenFile(fileName+".dat", os.O_RDWR|os.O_CREATE, 0644)
	if e != nil {
		if !os.IsPermission(e) {
			return fmt.Errorf("cannot create Volume Data %s.dat: %s", fileName, e)
		}
		if v.dataFile, e = os.Open(fileName + ".dat"); e != nil {
			return fmt.Errorf("cannot open Volume Data %s.dat: %s", fileName, e)
		}
		log.Printf("opening " + fileName + ".dat in READONLY mode")
		v.readOnly = true
	}
	if v.ReplicaType == CopyNil {
		e = v.readSuperBlock()
	} else {
		e = v.maybeWriteSuperBlock()
	}
	if e == nil && alsoLoadIndex {
		var indexFile *os.File
		if v.readOnly {
			if indexFile, e = os.Open(fileName + ".idx"); e != nil && !os.IsNotExist(e) {
				return fmt.Errorf("cannot open index file %s.idx: %s", fileName, e)
			}
			if indexFile != nil {
				log.Printf("converting %s.idx to %s.cdb", fileName, fileName)
				if e = ConvertIndexToCdb(fileName+".cdb", indexFile); e != nil {
					log.Printf("error converting %s.idx to %s.cdb: %s", fileName, fileName)
				} else {
					indexFile.Close()
					os.Remove(indexFile.Name())
					indexFile = nil
				}
			}
			v.nm, e = OpenCdbMap(fileName + ".cdb")
			return e
		} else {
			indexFile, e = os.OpenFile(fileName+".idx", os.O_RDWR|os.O_CREATE, 0644)
			if e != nil {
				return fmt.Errorf("cannot create Volume Data %s.dat: %s", fileName, e)
			}
		}
		v.nm, e = LoadNeedleMap(indexFile)
	}
	return e
}
func (v *Volume) Version() Version {
	return v.SuperBlock.Version
}
func (v *Volume) Size() int64 {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	stat, e := v.dataFile.Stat()
	if e == nil {
		return stat.Size()
	}
	fmt.Printf("Failed to read file size %s %s\n", v.dataFile.Name(), e.Error())
	return -1
}
func (v *Volume) Close() {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	v.nm.Close()
	_ = v.dataFile.Close()
}
func (v *Volume) maybeWriteSuperBlock() error {
	stat, e := v.dataFile.Stat()
	if e != nil {
		fmt.Printf("failed to stat datafile %s: %s", v.dataFile, e)
		return e
	}
	if stat.Size() == 0 {
		v.SuperBlock.Version = CurrentVersion
		_, e = v.dataFile.Write(v.SuperBlock.Bytes())
		if e != nil && os.IsPermission(e) {
			//read-only, but zero length - recreate it!
			if v.dataFile, e = os.Create(v.dataFile.Name()); e == nil {
				if _, e = v.dataFile.Write(v.SuperBlock.Bytes()); e == nil {
					v.readOnly = false
				}
			}
		}
	}
	return e
}
func (v *Volume) readSuperBlock() (err error) {
	if _, err = v.dataFile.Seek(0, 0); err != nil {
		return fmt.Errorf("cannot seek to the beginning of %s: %s", v.dataFile, err)
	}
	header := make([]byte, SuperBlockSize)
	if _, e := v.dataFile.Read(header); e != nil {
		return fmt.Errorf("cannot read superblock: %s", e)
	}
	v.SuperBlock, err = ParseSuperBlock(header)
	return err
}
func ParseSuperBlock(header []byte) (superBlock SuperBlock, err error) {
	superBlock.Version = Version(header[0])
	if superBlock.ReplicaType, err = NewReplicationTypeFromByte(header[1]); err != nil {
		err = fmt.Errorf("cannot read replica type: %s", err)
	}
	return
}
func (v *Volume) NeedToReplicate() bool {
	return v.ReplicaType.GetCopyCount() > 1
}

func (v *Volume) write(n *Needle) (size uint32, err error) {
	if v.readOnly {
		err = fmt.Errorf("%s is read-only", v.dataFile)
		return
	}
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	var offset int64
	if offset, err = v.dataFile.Seek(0, 2); err != nil {
		return
	}

	//ensure file writing starting from aligned positions
	if offset%NeedlePaddingSize != 0 {
		offset = offset + (NeedlePaddingSize - offset%NeedlePaddingSize)
		if offset, err = v.dataFile.Seek(offset, 0); err != nil {
			return
		}
	}

	if size, err = n.Append(v.dataFile, v.Version()); err != nil {
		if e := v.dataFile.Truncate(offset); e != nil {
			err = fmt.Errorf("%s\ncannot truncate %s: %s", err, v.dataFile, e)
		}
		return
	}
	nv, ok := v.nm.Get(n.Id)
	if !ok || int64(nv.Offset)*NeedlePaddingSize < offset {
		_, err = v.nm.Put(n.Id, uint32(offset/NeedlePaddingSize), n.Size)
	}
	return
}

func (v *Volume) delete(n *Needle) (uint32, error) {
	if v.readOnly {
		return 0, fmt.Errorf("%s is read-only", v.dataFile)
	}
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	nv, ok := v.nm.Get(n.Id)
	//fmt.Println("key", n.Id, "volume offset", nv.Offset, "data_size", n.Size, "cached size", nv.Size)
	if ok {
		var err error
		if err = v.nm.Delete(n.Id); err != nil {
			return nv.Size, err
		}
		if _, err = v.dataFile.Seek(int64(nv.Offset*NeedlePaddingSize), 0); err != nil {
			return nv.Size, err
		}
		_, err = n.Append(v.dataFile, v.Version())
		return nv.Size, err
	}
	return 0, nil
}

func (v *Volume) read(n *Needle) (int, error) {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	nv, ok := v.nm.Get(n.Id)
	if ok && nv.Offset > 0 {
		if _, err := v.dataFile.Seek(int64(nv.Offset)*NeedlePaddingSize, 0); err != nil {
			return -1, err
		}
		return n.Read(v.dataFile, nv.Size, v.Version())
	}
	return -1, errors.New("Not Found")
}

func (v *Volume) garbageLevel() float64 {
	return float64(v.nm.DeletedSize()) / float64(v.ContentSize())
}

func (v *Volume) compact() error {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()

	filePath := path.Join(v.dir, v.Id.String())
	return v.copyDataAndGenerateIndexFile(filePath+".cpd", filePath+".cpx")
}
func (v *Volume) commitCompact() error {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	_ = v.dataFile.Close()
	var e error
	if e = os.Rename(path.Join(v.dir, v.Id.String()+".cpd"), path.Join(v.dir, v.Id.String()+".dat")); e != nil {
		return e
	}
	if e = os.Rename(path.Join(v.dir, v.Id.String()+".cpx"), path.Join(v.dir, v.Id.String()+".idx")); e != nil {
		return e
	}
	if e = v.load(true); e != nil {
		return e
	}
	return nil
}
func (v *Volume) freeze() error {
	if v.readOnly {
		return nil
	}
	nm, ok := v.nm.(*NeedleMap)
	if !ok {
		return nil
	}
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	bn, _ := nakeFilename(v.dataFile.Name())
	cdbFn := bn + ".cdb"
	log.Printf("converting %s to %s", nm.indexFile.Name(), cdbFn)
	err := DumpNeedleMapToCdb(cdbFn, nm)
	if err != nil {
		return err
	}
	if v.nm, err = OpenCdbMap(cdbFn); err != nil {
		return err
	}
	nm.indexFile.Close()
	os.Remove(nm.indexFile.Name())
	v.readOnly = true
	return nil
}

func ScanVolumeFile(dirname string, id VolumeId,
	visitSuperBlock func(SuperBlock) error,
	visitNeedle func(n *Needle, offset uint32) error) (err error) {
	var v *Volume
	if v, err = LoadVolumeOnly(dirname, id); err != nil {
		return
	}
	if err = visitSuperBlock(v.SuperBlock); err != nil {
		return
	}

	version := v.Version()

	offset := uint32(SuperBlockSize)
	n, rest, e := ReadNeedleHeader(v.dataFile, version)
	if e != nil {
		err = fmt.Errorf("cannot read needle header: %s", e)
		return
	}
	for n != nil {
		if err = n.ReadNeedleBody(v.dataFile, version, rest); err != nil {
			err = fmt.Errorf("cannot read needle body: %s", err)
			return
		}
		if err = visitNeedle(n, offset); err != nil {
			return
		}
		offset += NeedleHeaderSize + rest
		if n, rest, err = ReadNeedleHeader(v.dataFile, version); err != nil {
			if err == io.EOF {
				return nil
			}
			return fmt.Errorf("cannot read needle header: %s", err)
		}
	}

	return
}

func (v *Volume) copyDataAndGenerateIndexFile(dstName, idxName string) (err error) {
	var (
		dst, idx *os.File
	)
	if dst, err = os.OpenFile(dstName, os.O_WRONLY|os.O_CREATE, 0644); err != nil {
		return
	}
	defer dst.Close()

	if idx, err = os.OpenFile(idxName, os.O_WRONLY|os.O_CREATE, 0644); err != nil {
		return
	}
	defer idx.Close()

	nm := NewNeedleMap(idx)
	new_offset := uint32(SuperBlockSize)

	err = ScanVolumeFile(v.dir, v.Id, func(superBlock SuperBlock) error {
		_, err = dst.Write(superBlock.Bytes())
		return err
	}, func(n *Needle, offset uint32) error {
		nv, ok := v.nm.Get(n.Id)
		//log.Println("file size is", n.Size, "rest", rest)
		if ok && nv.Offset*NeedlePaddingSize == offset {
			if nv.Size > 0 {
				if _, err = nm.Put(n.Id, new_offset/NeedlePaddingSize, n.Size); err != nil {
					return fmt.Errorf("cannot put needle: %s", err)
				}
				if _, err = n.Append(dst, v.Version()); err != nil {
					return fmt.Errorf("cannot append needle: %s", err)
				}
				new_offset += n.DiskSize()
				//log.Println("saving key", n.Id, "volume offset", old_offset, "=>", new_offset, "data_size", n.Size, "rest", rest)
			}
		}
		return nil
	})

	return
}
func (v *Volume) ContentSize() uint64 {
	return v.nm.ContentSize()
}
<|endoftext|>


<commit_before>package handler

import (
	"github.com/fuxiaohei/GoBlog/app/model"
	"github.com/fuxiaohei/GoBlog/app/utils"
	"github.com/fuxiaohei/GoInk"
	"net/url"
	"strconv"
	"strings"
)

func Login(context *GoInk.Context) {
	if context.Method == "POST" {
		data := context.Input()
		user := model.GetUserByName(data["user"])
		if user == nil {
			Json(context, false).End()
			return
		}
		if !user.CheckPassword(data["password"]) {
			Json(context, false).End()
			return
		}
		exp := 3600 * 24 * 3
		expStr := strconv.Itoa(exp)
		s := model.CreateToken(user, context, int64(exp))
		context.Cookie("token-user", strconv.Itoa(s.UserId), expStr)
		context.Cookie("token-value", s.Value, expStr)
		Json(context, true).End()
		return
	}
	if context.Cookie("token-value") != "" {
		context.Redirect("/admin/")
		return
	}
	context.Render("admin/login", nil)
}

func Auth(context *GoInk.Context) {
	tokenValue := context.Cookie("token-value")
	token := model.GetTokenByValue(tokenValue)
	if token == nil {
		context.Redirect("/logout/")
		context.End()
		return
	}
	if !token.IsValid() {
		context.Redirect("/logout/")
		context.End()
		return
	}
}

func Logout(context *GoInk.Context) {
	context.Cookie("token-user", "", "-3600")
	context.Cookie("token-value", "", "-3600")
	context.Redirect("/login/")
}

func TagArticles(ctx *GoInk.Context) {
	ctx.Layout("home")
	page, _ := strconv.Atoi(ctx.Param("page"))
	tag, _ := url.QueryUnescape(ctx.Param("tag"))
	size, _ := strconv.Atoi(model.GetSetting("article_size"))
	articles, pager := model.GetTaggedArticleList(tag, page, size)
	// fix dotted tag
	if len(articles) < 1 && strings.Contains(tag, "-") {
		articles, pager = model.GetTaggedArticleList(strings.Replace(tag, "-", ".", -1), page, size)
	}
	Theme(ctx).Layout("home").Render("index", map[string]interface{}{
		"Articles":    articles,
		"Pager":       pager,
		"SidebarHtml": SidebarHtml(ctx),
		"Tag":         tag,
		"Title":       tag,
	})
}

func Home(context *GoInk.Context) {
	context.Layout("home")
	page, _ := strconv.Atoi(context.Param("page"))
	size, _ := strconv.Atoi(model.GetSetting("article_size"))
	articles, pager := model.GetPublishArticleList(page, size)
	Theme(context).Layout("home").Render("index", map[string]interface{}{
		"Articles":    articles,
		"Pager":       pager,
		"SidebarHtml": SidebarHtml(context),
	})
}

func Article(context *GoInk.Context) {
	id, _ := strconv.Atoi(context.Param("id"))
	slug := context.Param("slug")
	article := model.GetContentById(id)
	if article == nil {
		context.Redirect("/")
		return
	}
	if article.Slug != slug || article.Type != "article" {
		context.Redirect("/")
		return
	}
	article.Hits++
	Theme(context).Layout("home").Render("article", map[string]interface{}{
		"Title":       article.Title,
		"Article":     article,
		"CommentHtml": CommentHtml(context, article),
	})
}

func Page(context *GoInk.Context) {
	id, _ := strconv.Atoi(context.Param("id"))
	slug := context.Param("slug")
	article := model.GetContentById(id)
	if article == nil || article.Status != "publish" {
		context.Redirect("/")
		return
	}
	if article.Slug != slug || article.Type != "page" {
		context.Redirect("/")
		return
	}
	article.Hits++
	Theme(context).Layout("home").Render("page", map[string]interface{}{
		"Title": article.Title,
		"Page":  article,
		//"CommentHtml": Comments(context, article),
	})
}

func TopPage(context *GoInk.Context) {
	slug := context.Param("slug")
	page := model.GetContentBySlug(slug)
	if page == nil || page.Status != "publish" {
		context.Redirect("/")
		return
	}
	if page.IsLinked && page.Type == "page" {
		Theme(context).Layout("home").Render("page", map[string]interface{}{
			"Title": page.Title,
			"Page":  page,
		})
		page.Hits++
		return
	}
	context.Redirect("/")
}

func Comment(context *GoInk.Context) {
	cid, _ := strconv.Atoi(context.Param("id"))
	if cid < 1 {
		Json(context, false).End()
		return
	}
	if model.GetContentById(cid) == nil {
		Json(context, false).End()
		return
	}
	data := context.Input()
	msg := validateComment(data)
	if msg != "" {
		Json(context, false).Set("msg", msg).End()
		return
	}
	co := new(model.Comment)
	co.Author = data["user"]
	co.Email = data["email"]
	co.Url = data["url"]
	co.Content = data["content"]
	co.Avatar = utils.Gravatar(co.Email, "50")
	co.Pid, _ = strconv.Atoi(data["pid"])
	co.Ip = context.Ip
	co.UserAgent = context.UserAgent
	co.IsAdmin = false
	model.CreateComment(cid, co)
	Json(context, true).Set("comment", co.ToJson()).End()
	model.CreateMessage("comment", co)
	context.Do("comment_created", co)
}

func validateComment(data map[string]string) string {
	if utils.IsEmptyString(data["user"]) || utils.IsEmptyString(data["content"]) {
		return "称呼，邮箱，内容必填"
	}
	if !utils.IsEmail(data["email"]) {
		return "邮箱格式错误"
	}
	if !utils.IsEmptyString(data["url"]) && !utils.IsURL(data["url"]) {
		return "网址格式错误"
	}
	return ""
}
<commit_msg>update home handler<commit_after>package handler

import (
	"github.com/fuxiaohei/GoBlog/app/model"
	"github.com/fuxiaohei/GoBlog/app/utils"
	"github.com/fuxiaohei/GoInk"
	"net/url"
	"strconv"
	"strings"
)

func Login(context *GoInk.Context) {
	if context.Method == "POST" {
		data := context.Input()
		user := model.GetUserByName(data["user"])
		if user == nil {
			Json(context, false).End()
			return
		}
		if !user.CheckPassword(data["password"]) {
			Json(context, false).End()
			return
		}
		exp := 3600 * 24 * 3
		expStr := strconv.Itoa(exp)
		s := model.CreateToken(user, context, int64(exp))
		context.Cookie("token-user", strconv.Itoa(s.UserId), expStr)
		context.Cookie("token-value", s.Value, expStr)
		Json(context, true).End()
		return
	}
	if context.Cookie("token-value") != "" {
		context.Redirect("/admin/")
		return
	}
	context.Render("admin/login", nil)
}

func Auth(context *GoInk.Context) {
	tokenValue := context.Cookie("token-value")
	token := model.GetTokenByValue(tokenValue)
	if token == nil {
		context.Redirect("/logout/")
		context.End()
		return
	}
	if !token.IsValid() {
		context.Redirect("/logout/")
		context.End()
		return
	}
}

func Logout(context *GoInk.Context) {
	context.Cookie("token-user", "", "-3600")
	context.Cookie("token-value", "", "-3600")
	context.Redirect("/login/")
}

func TagArticles(ctx *GoInk.Context) {
	ctx.Layout("home")
	page, _ := strconv.Atoi(ctx.Param("page"))
	tag, _ := url.QueryUnescape(ctx.Param("tag"))
	size := getArticleListSize()
	articles, pager := model.GetTaggedArticleList(tag, page, getArticleListSize())
	// fix dotted tag
	if len(articles) < 1 && strings.Contains(tag, "-") {
		articles, pager = model.GetTaggedArticleList(strings.Replace(tag, "-", ".", -1), page, size)
	}
	Theme(ctx).Layout("home").Render("index", map[string]interface{}{
		"Articles":    articles,
		"Pager":       pager,
		"SidebarHtml": SidebarHtml(ctx),
		"Tag":         tag,
		"Title":       tag,
	})
}

func Home(context *GoInk.Context) {
	context.Layout("home")
	page, _ := strconv.Atoi(context.Param("page"))
	articles, pager := model.GetPublishArticleList(page, getArticleListSize())
	data := map[string]interface{}{
		"Articles":    articles,
		"Pager":       pager,
		"SidebarHtml": SidebarHtml(context),
	}
	if page > 1 {
		data["Title"] = "第 " + strconv.Itoa(page) + " 页"
	}
	Theme(context).Layout("home").Render("index", data)
}

func Article(context *GoInk.Context) {
	id, _ := strconv.Atoi(context.Param("id"))
	slug := context.Param("slug")
	article := model.GetContentById(id)
	if article == nil {
		context.Redirect("/")
		return
	}
	if article.Slug != slug || article.Type != "article" {
		context.Redirect("/")
		return
	}
	article.Hits++
	Theme(context).Layout("home").Render("article", map[string]interface{}{
		"Title":       article.Title,
		"Article":     article,
		"CommentHtml": CommentHtml(context, article),
	})
}

func Page(context *GoInk.Context) {
	id, _ := strconv.Atoi(context.Param("id"))
	slug := context.Param("slug")
	article := model.GetContentById(id)
	if article == nil || article.Status != "publish" {
		context.Redirect("/")
		return
	}
	if article.Slug != slug || article.Type != "page" {
		context.Redirect("/")
		return
	}
	article.Hits++
	Theme(context).Layout("home").Render("page", map[string]interface{}{
		"Title": article.Title,
		"Page":  article,
		//"CommentHtml": Comments(context, article),
	})
}

func TopPage(context *GoInk.Context) {
	slug := context.Param("slug")
	page := model.GetContentBySlug(slug)
	if page == nil || page.Status != "publish" {
		context.Redirect("/")
		return
	}
	if page.IsLinked && page.Type == "page" {
		Theme(context).Layout("home").Render("page", map[string]interface{}{
			"Title": page.Title,
			"Page":  page,
		})
		page.Hits++
		return
	}
	context.Redirect("/")
}

func Comment(context *GoInk.Context) {
	cid, _ := strconv.Atoi(context.Param("id"))
	if cid < 1 {
		Json(context, false).End()
		return
	}
	if model.GetContentById(cid) == nil {
		Json(context, false).End()
		return
	}
	data := context.Input()
	msg := validateComment(data)
	if msg != "" {
		Json(context, false).Set("msg", msg).End()
		return
	}
	co := new(model.Comment)
	co.Author = data["user"]
	co.Email = data["email"]
	co.Url = data["url"]
	co.Content = data["content"]
	co.Avatar = utils.Gravatar(co.Email, "50")
	co.Pid, _ = strconv.Atoi(data["pid"])
	co.Ip = context.Ip
	co.UserAgent = context.UserAgent
	co.IsAdmin = false
	model.CreateComment(cid, co)
	Json(context, true).Set("comment", co.ToJson()).End()
	model.CreateMessage("comment", co)
	context.Do("comment_created", co)
}

func validateComment(data map[string]string) string {
	if utils.IsEmptyString(data["user"]) || utils.IsEmptyString(data["content"]) {
		return "称呼，邮箱，内容必填"
	}
	if !utils.IsEmail(data["email"]) {
		return "邮箱格式错误"
	}
	if !utils.IsEmptyString(data["url"]) && !utils.IsURL(data["url"]) {
		return "网址格式错误"
	}
	return ""
}

func getArticleListSize() int {
	size, _ := strconv.Atoi(model.GetSetting("article_size"))
	if size < 1 {
		size = 5
	}
	return size
}
<|endoftext|>


<commit_before>package sensorsanalytics

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
)

// Consumer sensors analytics consumer data
type Consumer interface {
	Send(message map[string]interface{}) error
	Flush() error
	Close() error
}

// DefaultConsumer 默认的 Consumer实现，逐条、同步的发送数据给接收服务器。
type DefaultConsumer struct {
	urlPrefix string
}

// NewDefaultConsumer 创建新的默认 Consumer
// :param serverURL: 服务器的 URL 地址。
func NewDefaultConsumer(serverURL string) *DefaultConsumer {
	var dc DefaultConsumer
	dc.urlPrefix = serverURL
	return &dc
}

// Send 发送数据
func (v *DefaultConsumer) Send(msg map[string]interface{}) error {
	type request struct {
		Data string `json:"data"`
		Gzip int    `json:"gzip"`
	}
	data, err := v.encodeMsg(msg)
	if err != nil {
		return fmt.Errorf("%s: %s", ErrIllegalDataException, err)
	}
	req := request{
		Data: data,
		Gzip: 0,
	}
	b := new(bytes.Buffer)
	json.NewEncoder(b).Encode(req)
	resp, err := http.Post(v.urlPrefix, "application/json; charset=utf-8", b)
	if err != nil {
		return fmt.Errorf("%s: %s", ErrNetworkException, err)
	}
	if resp.StatusCode != 200 && resp.StatusCode != 201 {
		return fmt.Errorf("%s: %s", ErrNetworkException, fmt.Sprintf("Error response status code [code=%d]", resp.StatusCode))
	}
	return nil
}

// Flush flush data
func (v *DefaultConsumer) Flush() error {
	return nil
}

// Close close consumer
func (v *DefaultConsumer) Close() error {
	return nil
}

func (v *DefaultConsumer) encodeMsg(msg map[string]interface{}) (string, error) {
	s, err := json.Marshal(msg)
	if err != nil {
		return "", err
	}
	data := base64.StdEncoding.EncodeToString(s)
	return data, nil
}
<commit_msg>Add batch consumer<commit_after>package sensorsanalytics

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
)

// Consumer sensors analytics consumer data
type Consumer interface {
	Send(message map[string]interface{}) error
	Flush() error
	Close() error
}

// SARequest sensorsdata anaalytics request
type SARequest struct {
	Data string `json:"data"`
	Gzip int    `json:"gzip"`
}

// DefaultConsumer 默认的 Consumer实现，逐条、同步的发送数据给接收服务器。
type DefaultConsumer struct {
	urlPrefix string
}

// NewDefaultConsumer 创建新的默认 Consumer
// :param serverURL: 服务器的 URL 地址。
func NewDefaultConsumer(serverURL string) *DefaultConsumer {
	var c DefaultConsumer
	c.urlPrefix = serverURL
	return &c
}

// Send 发送数据
func (c *DefaultConsumer) Send(msg map[string]interface{}) error {
	data, err := c.encodeMsg(msg)
	if err != nil {
		return fmt.Errorf("%s: %s", ErrIllegalDataException, err)
	}
	req := SARequest{
		Data: data,
		Gzip: 0,
	}
	b := new(bytes.Buffer)
	json.NewEncoder(b).Encode(req)
	resp, err := http.Post(c.urlPrefix, "application/json; charset=utf-8", b)
	if err != nil {
		return fmt.Errorf("%s: %s", ErrNetworkException, err)
	}
	if resp.StatusCode != 200 && resp.StatusCode != 201 {
		return fmt.Errorf("%s: %s", ErrNetworkException, fmt.Sprintf("Error response status code [code=%d]", resp.StatusCode))
	}
	return nil
}

// Flush flush data
func (c *DefaultConsumer) Flush() error {
	return nil
}

// Close close consumer
func (c *DefaultConsumer) Close() error {
	return nil
}

func (c *DefaultConsumer) encodeMsg(msg map[string]interface{}) (string, error) {
	s, err := json.Marshal(msg)
	if err != nil {
		return "", err
	}
	data := base64.StdEncoding.EncodeToString(s)
	return data, nil
}

// BatchConsumer  批量发送数据的 Consumer，当且仅当数据达到 buffer_size 参数指定的量时，才将数据进行发送。
type BatchConsumer struct {
	DefaultConsumer
	bufferSize int
	buffer     []string
}

// NewBatchConsumer 创建新的 batch consumer
func NewBatchConsumer(serverURL string, bufferSize int) *BatchConsumer {
	var c BatchConsumer
	c.urlPrefix = serverURL
	if bufferSize > 0 && bufferSize <= 50 {
		c.bufferSize = bufferSize
	} else {
		c.bufferSize = 20
	}
	buffer := make([]string, c.bufferSize)
	c.buffer = buffer
	return &c
}

// Send 新的 msg 加入 buffer
func (c *BatchConsumer) Send(msg map[string]interface{}) error {
	data, err := c.encodeMsg(msg)
	if err != nil {
		return fmt.Errorf("%s: %s", ErrIllegalDataException, err)
	}
	c.buffer = append(c.buffer, data)
	if len(c.buffer) >= c.bufferSize {
		return c.Flush()
	}
	return nil
}

// Flush  用户可以主动调用 flush 接口，以便在需要的时候立即进行数据发送。
func (c *BatchConsumer) Flush() error {
	for _, v := range c.buffer {
		req := SARequest{
			Data: v,
			Gzip: 0,
		}
		b := new(bytes.Buffer)
		json.NewEncoder(b).Encode(req)
		resp, err := http.Post(c.urlPrefix, "application/json; charset=utf-8", b)
		if err != nil {
			log.Printf("%s: %s", ErrNetworkException, err)
		}
		if resp.StatusCode != 200 && resp.StatusCode != 201 {
			log.Printf("%s: %s", ErrNetworkException, fmt.Sprintf("Error response status code [code=%d]", resp.StatusCode))
		}
	}
	c.buffer = make([]string, c.bufferSize)
	return nil
}

// Close 在发送完成时，调用此接口以保证数据发送完成。
func (c *BatchConsumer) Close() error {
	return c.Flush()
}
<|endoftext|>


<commit_before>package http

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

type consumer struct {
	conn   net.Conn
	es     *eventSource
	in     chan []byte
	staled bool
}

func newConsumer(resp http.ResponseWriter, req *http.Request, es *eventSource) (*consumer, error) {
	conn, _, err := resp.(http.Hijacker).Hijack()
	if err != nil {
		return nil, err
	}

	consumer := &consumer{
		conn:   conn,
		es:     es,
		in:     make(chan []byte, 10),
		staled: false,
	}

	headers := [][]byte{
		[]byte("HTTP/1.1 200 OK"),
		[]byte("Content-Type: text/event-stream"),
	}

	for i := 0; i < len(headers); i++ {
		conn.Write(headers[i])
		conn.Write([]byte("\n"))
	}

	if es.customHeadersFunc != nil {
		customHeaders := es.customHeadersFunc(req)
		for i := 0; i < len(customHeaders); i++ {
			conn.Write(customHeaders[i])
			conn.Write([]byte("\n"))
		}
	}

	_, err = conn.Write([]byte(fmt.Sprintf("retry: %d\n\n", es.retry/time.Millisecond)))

	if err != nil {
		conn.Close()
		return nil, err
	}

	go func() {
		idleTimer := time.NewTimer(es.idleTimeout)
		defer idleTimer.Stop()
		for {
			select {
			case message, open := <-consumer.in:
				if !open {
					consumer.conn.Close()
					return
				}
				conn.SetWriteDeadline(time.Now().Add(consumer.es.timeout))
				_, err := conn.Write(message)
				if err != nil {
					netErr, ok := err.(net.Error)
					if !ok || !netErr.Timeout() || consumer.es.closeOnTimeout {
						consumer.staled = true
						consumer.conn.Close()
						consumer.es.staled <- consumer
						return
					}
				}
				idleTimer.Reset(es.idleTimeout)
			case <-idleTimer.C:
				consumer.conn.Close()
				consumer.es.staled <- consumer
				return
			}
		}
	}()

	return consumer, nil
}
<commit_msg>use range instead of for loop<commit_after>package http

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

type consumer struct {
	conn   net.Conn
	es     *eventSource
	in     chan []byte
	staled bool
}

func newConsumer(resp http.ResponseWriter, req *http.Request, es *eventSource) (*consumer, error) {
	conn, _, err := resp.(http.Hijacker).Hijack()
	if err != nil {
		return nil, err
	}

	consumer := &consumer{
		conn:   conn,
		es:     es,
		in:     make(chan []byte, 10),
		staled: false,
	}

	headers := [][]byte{
		[]byte("HTTP/1.1 200 OK"),
		[]byte("Content-Type: text/event-stream"),
	}

	for _, header := range headers {
		conn.Write(header)
		conn.Write([]byte("\n"))
	}

	if es.customHeadersFunc != nil {
		for _, header := range es.customHeadersFunc(req) {
			conn.Write(header)
			conn.Write([]byte("\n"))
		}
	}

	_, err = conn.Write([]byte(fmt.Sprintf("retry: %d\n\n", es.retry/time.Millisecond)))

	if err != nil {
		conn.Close()
		return nil, err
	}

	go func() {
		idleTimer := time.NewTimer(es.idleTimeout)
		defer idleTimer.Stop()
		for {
			select {
			case message, open := <-consumer.in:
				if !open {
					consumer.conn.Close()
					return
				}
				conn.SetWriteDeadline(time.Now().Add(consumer.es.timeout))
				_, err := conn.Write(message)
				if err != nil {
					netErr, ok := err.(net.Error)
					if !ok || !netErr.Timeout() || consumer.es.closeOnTimeout {
						consumer.staled = true
						consumer.conn.Close()
						consumer.es.staled <- consumer
						return
					}
				}
				idleTimer.Reset(es.idleTimeout)
			case <-idleTimer.C:
				consumer.conn.Close()
				consumer.es.staled <- consumer
				return
			}
		}
	}()

	return consumer, nil
}
<|endoftext|>


<commit_before>package widgets

import (
	"fmt"
	"io"

	"github.com/jroimartin/gocui"
)

type Items interface {
	Display(w io.Writer)
	Len() int
}

type EmptyItems struct{}

func (e *EmptyItems) Display(w io.Writer) {}
func (e *EmptyItems) Len() int            { return 0 }

type SelectionEvent int
type Selected struct {
	Type  SelectionEvent
	Index int
	Items Items
}

const (
	OnSelect SelectionEvent = iota
	OnEnter
)

type list struct {
	*UI
	items    Items
	current  int
	selected chan *Selected
}

func NewList(ui *UI) (*list, chan *Selected) {
	selected := make(chan *Selected)
	l := &list{
		ui,
		&EmptyItems{},
		-1,
		selected,
	}

	l.AddLocalKey(gocui.KeyArrowUp, "Move selection up one line", l.Previous)
	l.AddLocalKey('k', "Move selection up one line", l.Previous)
	l.AddLocalKey(gocui.KeyArrowDown, "Move selection down one line", l.Next)
	l.AddLocalKey('j', "Move selection down one line", l.Next)

	l.AddLocalKey(gocui.KeyEnter, "Select current line", func() { l.fire(OnEnter) })

	return l, selected
}

func (l *list) SetSelection(index int) {
	l.Update(func(v *gocui.View) {
		count := l.items.Len()

		if count == 0 {
			return
		}

		if index < 0 || index >= count {
			fmt.Print("\a")
			return
		}

		if l.current == -1 {
			l.current = index
			v.SetOrigin(0, l.current)
		} else {
			moveDistance := index - l.current
			l.current = index
			v.MoveCursor(0, moveDistance, false)
		}
		l.fire(OnSelect)
	})
}

func (l *list) fire(event SelectionEvent) {
	go func() {
		l.selected <- &Selected{event, l.current, l.items}
	}()
}

func (l *list) SetItems(items Items, index int) {
	l.Update(func(v *gocui.View) {
		v.Clear()
		l.items = items

		items.Display(v)

		l.SetSelection(index)
	})
}

func (l *list) Next()     { l.SetSelection(l.current + 1) }
func (l *list) Previous() { l.SetSelection(l.current - 1) }
<commit_msg>Don't allow list selection before the list has been populated.<commit_after>package widgets

import (
	"fmt"
	"io"

	"github.com/jroimartin/gocui"
)

type Items interface {
	Display(w io.Writer)
	Len() int
}

type EmptyItems struct{}

func (e *EmptyItems) Display(w io.Writer) {}
func (e *EmptyItems) Len() int            { return 0 }

type SelectionEvent int
type Selected struct {
	Type  SelectionEvent
	Index int
	Items Items
}

const (
	OnSelect SelectionEvent = iota
	OnEnter
)

type list struct {
	*UI
	items    Items
	current  int
	selected chan *Selected
}

func NewList(ui *UI) (*list, chan *Selected) {
	selected := make(chan *Selected)
	l := &list{
		ui,
		&EmptyItems{},
		-1,
		selected,
	}

	l.AddLocalKey(gocui.KeyArrowUp, "Move selection up one line", l.Previous)
	l.AddLocalKey('k', "Move selection up one line", l.Previous)
	l.AddLocalKey(gocui.KeyArrowDown, "Move selection down one line", l.Next)
	l.AddLocalKey('j', "Move selection down one line", l.Next)

	l.AddLocalKey(gocui.KeyEnter, "Select current line", func() { l.fire(OnEnter) })

	return l, selected
}

func (l *list) SetSelection(index int) {
	l.Update(func(v *gocui.View) {
		count := l.items.Len()

		if count == 0 {
			return
		}

		if index < 0 || index >= count {
			fmt.Print("\a")
			return
		}

		if l.current == -1 {
			l.current = index
			v.SetOrigin(0, l.current)
		} else {
			moveDistance := index - l.current
			l.current = index
			v.MoveCursor(0, moveDistance, false)
		}
		l.fire(OnSelect)
	})
}

func (l *list) fire(event SelectionEvent) {
	go func() {
		if l.current != -1 { // keep users from selecting items before they have been populated
			l.selected <- &Selected{event, l.current, l.items}
		}
	}()
}

func (l *list) SetItems(items Items, index int) {
	l.Update(func(v *gocui.View) {
		v.Clear()
		l.items = items

		items.Display(v)

		l.SetSelection(index)
	})
}

func (l *list) Next()     { l.SetSelection(l.current + 1) }
func (l *list) Previous() { l.SetSelection(l.current - 1) }
<|endoftext|>


<commit_before>// NOTE: THIS API IS UNSTABLE AND WILL BE MOVED TO ITS OWN PACKAGE OR REFACTORED
// OUT.

package neutrino

import "github.com/btcsuite/btcd/wire"

// messageType describes the type of blockMessage.
type messageType int

const (
	connectBasic messageType = iota
	connectExt
	disconnect
)

// blockMessage is a notification from the block manager to a block
// subscription's goroutine to be forwarded on via the appropriate channel.
type blockMessage struct {
	header  *wire.BlockHeader
	msgType messageType
}

// blockSubscription allows a client to subscribe to and unsubscribe from block
// connect and disconnect notifications.
// TODO(aakselrod): Move this to its own package so that the subscriber can't
// access internals, in particular the notifyBlock and intQuit members.
type blockSubscription struct {
	onConnectBasic chan<- wire.BlockHeader
	onConnectExt   chan<- wire.BlockHeader
	onDisconnect   chan<- wire.BlockHeader
	quit           <-chan struct{}

	notifyBlock chan *blockMessage
	intQuit     chan struct{}
}

// sendSubscribedMsg sends all block subscribers a message if they request this
// type.
// TODO(aakselrod): Refactor so we're able to handle more message types in new
// package.
func (s *ChainService) sendSubscribedMsg(bm *blockMessage) {
	var subChan chan<- wire.BlockHeader
	s.mtxSubscribers.RLock()
	for sub := range s.blockSubscribers {
		switch bm.msgType {
		case connectBasic:
			subChan = sub.onConnectBasic
		case connectExt:
			subChan = sub.onConnectExt
		case disconnect:
			subChan = sub.onDisconnect
		default:
			// TODO: Return a useful error when factored out into
			// its own package.
			panic("invalid message type")
		}
		if subChan != nil {
			select {
			case sub.notifyBlock <- bm:
			case <-sub.quit:
			case <-sub.intQuit:
			}
		}
	}
	s.mtxSubscribers.RUnlock()
}

// subscribeBlockMsg handles adding block subscriptions to the ChainService.
// TODO(aakselrod): move this to its own package and refactor so that we're
// not modifying an object held by the caller.
func (s *ChainService) subscribeBlockMsg(onConnectBasic, onConnectExt,
	onDisconnect chan<- wire.BlockHeader,
	quit <-chan struct{}) *blockSubscription {
	s.mtxSubscribers.Lock()
	defer s.mtxSubscribers.Unlock()
	subscription := blockSubscription{
		onConnectBasic: onConnectBasic,
		onConnectExt:   onConnectExt,
		onDisconnect:   onDisconnect,
		quit:           quit,
		notifyBlock:    make(chan *blockMessage),
		intQuit:        make(chan struct{}),
	}
	s.blockSubscribers[&subscription] = struct{}{}
	go subscription.subscriptionHandler()
	return &subscription
}

// unsubscribeBlockMsgs handles removing block subscriptions from the
// ChainService.
// TODO(aakselrod): move this to its own package and refactor so that we're
// not depending on the caller to not modify the argument between subscribe and
// unsubscribe.
func (s *ChainService) unsubscribeBlockMsgs(subscription *blockSubscription) {
	s.mtxSubscribers.Lock()
	delete(s.blockSubscribers, subscription)
	s.mtxSubscribers.Unlock()
	close(subscription.intQuit)

	// Drain the inbound notification channel
cleanup:
	for {
		select {
		case <-subscription.notifyBlock:
		default:
			break cleanup
		}
	}
}

// subscriptionHandler must be run as a goroutine and queues notification
// messages from the chain service to the subscriber.
func (s *blockSubscription) subscriptionHandler() {
	// Start with a small queue; it will grow if needed.
	ntfns := make([]*blockMessage, 0, 5)
	var next *blockMessage

	// Try to send on the specified channel. If a new message arrives while
	// we try to send, queue it and continue with the loop. If a quit signal
	// is sent, let the loop know.
	selectChan := func(notify chan<- wire.BlockHeader) bool {
		if notify == nil {
			select {
			case <-s.quit:
				return false
			case <-s.intQuit:
				return false
			default:
				return true
			}
		}
		select {
		case notify <- *next.header:
			next = nil
			return true
		case queueMsg := <-s.notifyBlock:
			ntfns = append(ntfns, queueMsg)
			return true
		case <-s.quit:
			return false
		case <-s.intQuit:
			return false
		}
	}

	// Loop until we get a signal on s.quit or s.intQuit.
	for {
		if next != nil {
			// If selectChan returns false, we were signalled on
			// s.quit or s.intQuit.
			switch next.msgType {
			case connectBasic:
				if !selectChan(s.onConnectBasic) {
					return
				}
			case connectExt:
				if !selectChan(s.onConnectExt) {
					return
				}
			case disconnect:
				if !selectChan(s.onDisconnect) {
					return
				}
			}
		} else {
			// Next notification is nil, so see if we can get a
			// notification from the queue. If not, we wait for a
			// notification on s.notifyBlock or quit if signalled.
			if len(ntfns) > 0 {
				next = ntfns[0]
				ntfns = ntfns[1:]
			} else {
				select {
				case next = <-s.notifyBlock:
				case <-s.quit:
					return
				case <-s.intQuit:
					return
				}
			}
		}
	}
}
<commit_msg>subscriptions: remove connectExt as the ext filter no longer exists<commit_after>// NOTE: THIS API IS UNSTABLE AND WILL BE MOVED TO ITS OWN PACKAGE OR REFACTORED
// OUT.

package neutrino

import "github.com/btcsuite/btcd/wire"

// messageType describes the type of blockMessage.
type messageType int

const (
	connectBasic messageType = iota
	disconnect
)

// blockMessage is a notification from the block manager to a block
// subscription's goroutine to be forwarded on via the appropriate channel.
type blockMessage struct {
	header  *wire.BlockHeader
	msgType messageType
}

// blockSubscription allows a client to subscribe to and unsubscribe from block
// connect and disconnect notifications.
// TODO(aakselrod): Move this to its own package so that the subscriber can't
// access internals, in particular the notifyBlock and intQuit members.
type blockSubscription struct {
	onConnectBasic chan<- wire.BlockHeader
	onDisconnect   chan<- wire.BlockHeader
	quit           <-chan struct{}

	notifyBlock chan *blockMessage
	intQuit     chan struct{}
}

// sendSubscribedMsg sends all block subscribers a message if they request this
// type.
// TODO(aakselrod): Refactor so we're able to handle more message types in new
// package.
func (s *ChainService) sendSubscribedMsg(bm *blockMessage) {
	var subChan chan<- wire.BlockHeader
	s.mtxSubscribers.RLock()
	for sub := range s.blockSubscribers {
		switch bm.msgType {
		case connectBasic:
			subChan = sub.onConnectBasic
		case disconnect:
			subChan = sub.onDisconnect
		default:
			// TODO: Return a useful error when factored out into
			// its own package.
			panic("invalid message type")
		}
		if subChan != nil {
			select {
			case sub.notifyBlock <- bm:
			case <-sub.quit:
			case <-sub.intQuit:
			}
		}
	}
	s.mtxSubscribers.RUnlock()
}

// subscribeBlockMsg handles adding block subscriptions to the ChainService.
// TODO(aakselrod): move this to its own package and refactor so that we're
// not modifying an object held by the caller.
func (s *ChainService) subscribeBlockMsg(onConnectBasic, onConnectExt,
	quit <-chan struct{}) *blockSubscription {
	s.mtxSubscribers.Lock()
	defer s.mtxSubscribers.Unlock()
	subscription := blockSubscription{
		onConnectBasic: onConnectBasic,
		onDisconnect:   onDisconnect,
		quit:           quit,
		notifyBlock:    make(chan *blockMessage),
		intQuit:        make(chan struct{}),
	}
	s.blockSubscribers[&subscription] = struct{}{}
	go subscription.subscriptionHandler()
	return &subscription
}

// unsubscribeBlockMsgs handles removing block subscriptions from the
// ChainService.
//
// TODO(aakselrod): move this to its own package and refactor so that we're
// not depending on the caller to not modify the argument between subscribe and
// unsubscribe.
func (s *ChainService) unsubscribeBlockMsgs(subscription *blockSubscription) {
	s.mtxSubscribers.Lock()
	delete(s.blockSubscribers, subscription)
	s.mtxSubscribers.Unlock()

	close(subscription.intQuit)

	// Drain the inbound notification channel
cleanup:
	for {
		select {
		case <-subscription.notifyBlock:
		default:
			break cleanup
		}
	}
}

// subscriptionHandler must be run as a goroutine and queues notification
// messages from the chain service to the subscriber.
func (s *blockSubscription) subscriptionHandler() {
	// Start with a small queue; it will grow if needed.
	ntfns := make([]*blockMessage, 0, 5)
	var next *blockMessage

	// Try to send on the specified channel. If a new message arrives while
	// we try to send, queue it and continue with the loop. If a quit
	// signal is sent, let the loop know.
	selectChan := func(notify chan<- wire.BlockHeader) bool {
		if notify == nil {
			select {
			case <-s.quit:
				return false
			case <-s.intQuit:
				return false
			default:
				return true
			}
		}
		select {
		case notify <- *next.header:
			next = nil
			return true
		case queueMsg := <-s.notifyBlock:
			ntfns = append(ntfns, queueMsg)
			return true
		case <-s.quit:
			return false
		case <-s.intQuit:
			return false
		}
	}

	// Loop until we get a signal on s.quit or s.intQuit.
	for {
		if next != nil {
			// If selectChan returns false, we were signalled on
			// s.quit or s.intQuit.
			switch next.msgType {
			case connectBasic:
				if !selectChan(s.onConnectBasic) {
					return
				}
			case disconnect:
				if !selectChan(s.onDisconnect) {
					return
				}
			}
		} else {
			// Next notification is nil, so see if we can get a
			// notification from the queue. If not, we wait for a
			// notification on s.notifyBlock or quit if signalled.
			if len(ntfns) > 0 {
				next = ntfns[0]
				ntfns = ntfns[1:]
			} else {
				select {
				case next = <-s.notifyBlock:
				case <-s.quit:
					return
				case <-s.intQuit:
					return
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>package pgsql

import (
	"database/sql"
	"fmt"

	"github.com/coreos/clair/database"
	"github.com/coreos/clair/utils"
	cerrors "github.com/coreos/clair/utils/errors"
	"github.com/coreos/clair/utils/types"
	"github.com/guregu/null/zero"
)

func (pgSQL *pgSQL) FindVulnerability(namespaceName, name string) (database.Vulnerability, error) {
	vulnerability := database.Vulnerability{
		Name: name,
		Namespace: database.Namespace{
			Name: namespaceName,
		},
	}

	// Find Vulnerability.
	rows, err := pgSQL.Query(getQuery("f_vulnerability"), namespaceName, name)
	if err != nil {
		return vulnerability, handleError("f_vulnerability", err)
	}
	defer rows.Close()

	// Iterate to scan the Vulnerability and its FixedIn FeatureVersions.
	for rows.Next() {
		var featureVersionID zero.Int
		var featureVersionVersion zero.String
		var featureVersionFeatureName zero.String

		err := rows.Scan(&vulnerability.ID, &vulnerability.Namespace.ID, &vulnerability.Description,
			&vulnerability.Link, &vulnerability.Severity, &featureVersionVersion, &featureVersionID,
			&featureVersionFeatureName)
		if err != nil {
			return vulnerability, handleError("f_vulnerability.Scan()", err)
		}

		if !featureVersionID.IsZero() {
			// Note that the ID we fill in featureVersion is actually a Feature ID, and not
			// a FeatureVersion ID.
			featureVersion := database.FeatureVersion{
				Model: database.Model{ID: int(featureVersionID.Int64)},
				Feature: database.Feature{
					Model:     database.Model{ID: int(featureVersionID.Int64)},
					Namespace: vulnerability.Namespace,
					Name:      featureVersionFeatureName.String,
				},
				Version: types.NewVersionUnsafe(featureVersionVersion.String),
			}
			vulnerability.FixedIn = append(vulnerability.FixedIn, featureVersion)
		}
	}
	if err = rows.Err(); err != nil {
		return vulnerability, handleError("s_featureversions_vulnerabilities.Rows()", err)
	}
	if vulnerability.ID == 0 {
		return vulnerability, cerrors.ErrNotFound
	}

	return vulnerability, nil
}

// FixedIn.Namespace are not necessary, they are overwritten by the vuln.
// By setting the fixed version to minVersion, we can say that the vuln does'nt affect anymore.
func (pgSQL *pgSQL) InsertVulnerabilities(vulnerabilities []database.Vulnerability) error {
	for _, vulnerability := range vulnerabilities {
		err := pgSQL.insertVulnerability(vulnerability)
		if err != nil {
			fmt.Printf("%#v\n", vulnerability)
			return err
		}
	}
	return nil
}

func (pgSQL *pgSQL) insertVulnerability(vulnerability database.Vulnerability) error {
	// Verify parameters
	if vulnerability.Name == "" || len(vulnerability.FixedIn) == 0 ||
		vulnerability.Namespace.Name == "" || !vulnerability.Severity.IsValid() {
		log.Warning("could not insert an invalid vulnerability")
		return cerrors.NewBadRequestError("could not insert an invalid vulnerability")
	}

	for _, fixedInFeatureVersion := range vulnerability.FixedIn {
		if fixedInFeatureVersion.Feature.Namespace.Name != "" &&
			fixedInFeatureVersion.Feature.Namespace.Name != vulnerability.Namespace.Name {
			msg := "could not insert an invalid vulnerability: FixedIn FeatureVersion must be in the " +
				"same namespace as the Vulnerability"
			log.Warning(msg)
			return cerrors.NewBadRequestError(msg)
		}
	}

	// Find or insert Vulnerability's Namespace.
	namespaceID, err := pgSQL.insertNamespace(vulnerability.Namespace)
	if err != nil {
		return err
	}

	// Find vulnerability and its Vulnerability_FixedIn_Features.
	existingVulnerability, err := pgSQL.FindVulnerability(vulnerability.Namespace.Name,
		vulnerability.Name)
	if err != nil && err != cerrors.ErrNotFound {
		return err
	}

	// Compute new/updated FixedIn FeatureVersions.
	var newFixedInFeatureVersions []database.FeatureVersion
	var updatedFixedInFeatureVersions []database.FeatureVersion
	if existingVulnerability.ID == 0 {
		newFixedInFeatureVersions = vulnerability.FixedIn
	} else {
		newFixedInFeatureVersions, updatedFixedInFeatureVersions = diffFixedIn(vulnerability,
			existingVulnerability)
	}

	if len(newFixedInFeatureVersions) == 0 && len(updatedFixedInFeatureVersions) == 0 {
		// Nothing to do.
		return nil
	}

	// Insert or find the new Feature.
	// We already have the Feature IDs in updatedFixedInFeatureVersions because diffFixedIn fills them
	// in using the existing vulnerability's FixedIn FeatureVersions. Note that even if FixedIn
	// is type FeatureVersion, the actual stored ID in these structs are the Feature IDs.
	//
	// Also, we enforce the namespace of the FeatureVersion in case it was empty. There is a test
	// above to ensure that the passed Namespace is either the same as the vulnerability or empty.
	//
	// TODO(Quentin-M): Batch me.
	for i := 0; i < len(newFixedInFeatureVersions); i++ {
		newFixedInFeatureVersions[i].Feature.Namespace.Name = vulnerability.Namespace.Name
		newFixedInFeatureVersions[i].Feature.ID, err = pgSQL.insertFeature(newFixedInFeatureVersions[i].Feature)
		if err != nil {
			return err
		}
	}

	// Begin transaction.
	tx, err := pgSQL.Begin()
	if err != nil {
		tx.Rollback()
		return handleError("insertVulnerability.Begin()", err)
	}

	// Set transaction as SERIALIZABLE.
	// This is how we ensure that the data in Vulnerability_Affects_FeatureVersion is always
	// consistent.
	_, err = tx.Exec(getQuery("set_tx_serializable"))
	if err != nil {
		tx.Rollback()
		return handleError("insertFeatureVersion.set_tx_serializable", err)
	}

	if existingVulnerability.ID == 0 {
		// Insert new vulnerability.
		err = tx.QueryRow(getQuery("i_vulnerability"), namespaceID, vulnerability.Name,
			vulnerability.Description, vulnerability.Link, &vulnerability.Severity).Scan(&vulnerability.ID)
		if err != nil {
			tx.Rollback()
			return handleError("i_vulnerability", err)
		}
	} else {
		// Update vulnerability
		_, err = tx.Exec(getQuery("u_vulnerability"), existingVulnerability.ID,
			vulnerability.Description, vulnerability.Link, &vulnerability.Severity)
		if err != nil {
			tx.Rollback()
			return handleError("u_vulnerability", err)
		}

		vulnerability.ID = existingVulnerability.ID
	}

	// Update Vulnerability_FixedIn_Feature and Vulnerability_Affects_FeatureVersion now.
	err = pgSQL.updateVulnerabilityFeatureVersions(tx, &vulnerability, &existingVulnerability, newFixedInFeatureVersions, updatedFixedInFeatureVersions)
	if err != nil {
		tx.Rollback()
		return err
	}

	// Commit transaction.
	err = tx.Commit()
	if err != nil {
		tx.Rollback()
		return handleError("insertVulnerability.Commit()", err)
	}

	return nil
}

func diffFixedIn(vulnerability, existingVulnerability database.Vulnerability) (newFixedIn, updatedFixedIn []database.FeatureVersion) {
	// Build FeatureVersion.Feature.Namespace.Name:FeatureVersion.Feature.Name (NaN) structures.
	vulnerabilityFixedInNameMap, vulnerabilityFixedInNameSlice := createFeatureVersionNameMap(vulnerability.FixedIn)
	existingFixedInMapNameMap, existingFixedInNameSlice := createFeatureVersionNameMap(existingVulnerability.FixedIn)

	// Calculate the new FixedIn FeatureVersion NaN and updated ones.
	newFixedInName := utils.CompareStringLists(vulnerabilityFixedInNameSlice,
		existingFixedInNameSlice)
	updatedFixedInName := utils.CompareStringListsInBoth(vulnerabilityFixedInNameSlice,
		existingFixedInNameSlice)

	for _, nan := range newFixedInName {
		newFixedIn = append(newFixedIn, vulnerabilityFixedInNameMap[nan])
	}
	for _, nan := range updatedFixedInName {
		fv := existingFixedInMapNameMap[nan]
		fv.Version = vulnerabilityFixedInNameMap[nan].Version
		updatedFixedIn = append(updatedFixedIn, fv)
	}

	return
}

func createFeatureVersionNameMap(features []database.FeatureVersion) (map[string]database.FeatureVersion, []string) {
	m := make(map[string]database.FeatureVersion, 0)
	s := make([]string, 0, len(features))

	for i := 0; i < len(features); i++ {
		featureVersion := features[i]
		m[featureVersion.Feature.Name] = featureVersion
		s = append(s, featureVersion.Feature.Name)
	}

	return m, s
}

func (pgSQL *pgSQL) updateVulnerabilityFeatureVersions(tx *sql.Tx, vulnerability, existingVulnerability *database.Vulnerability, newFixedInFeatureVersions, updatedFixedInFeatureVersions []database.FeatureVersion) error {
	var fixedInID int

	for _, fv := range newFixedInFeatureVersions {
		if fv.Version == types.MinVersion {
			// We don't want to mark a Feature as fixed in MinVersion. MinVersion only makes sense when a
			// Feature is already marked as fixed in some version, in which case we would be in the
			// "updatedFixedInFeatureVersions" loop and removes the fixed in mark.
			continue
		}

		// Insert Vulnerability_FixedIn_Feature.
		err := tx.QueryRow(getQuery("i_vulnerability_fixedin_feature"), vulnerability.ID, fv.Feature.ID,
			&fv.Version).Scan(&fixedInID)
		if err != nil {
			return handleError("i_vulnerability_fixedin_feature", err)
		}

		// Insert Vulnerability_Affects_FeatureVersion.
		err = linkVulnerabilityToFeatureVersions(tx, fixedInID, vulnerability.ID, fv.Feature.ID,
			fv.Version)
		if err != nil {
			return err
		}
	}

	for _, fv := range updatedFixedInFeatureVersions {
		if fv.Version != types.MinVersion {
			// Update Vulnerability_FixedIn_Feature.
			err := tx.QueryRow(getQuery("u_vulnerability_fixedin_feature"), vulnerability.ID,
				fv.Feature.ID, &fv.Version).Scan(&fixedInID)
			if err != nil {
				return handleError("u_vulnerability_fixedin_feature", err)
			}

			// Drop all old Vulnerability_Affects_FeatureVersion.
			_, err = tx.Exec(getQuery("r_vulnerability_affects_featureversion"), fixedInID)
			if err != nil {
				return handleError("r_vulnerability_affects_featureversion", err)
			}

			// Insert Vulnerability_Affects_FeatureVersion.
			err = linkVulnerabilityToFeatureVersions(tx, fixedInID, vulnerability.ID, fv.Feature.ID,
				fv.Version)
			if err != nil {
				return err
			}
		} else {
			// Updating FixedIn by saying that the fixed version is the lowest possible version, it
			// basically means that the vulnerability doesn't affect the feature (anymore).
			// Drop it from Vulnerability_FixedIn_Feature and Vulnerability_Affects_FeatureVersion.
			err := tx.QueryRow(getQuery("r_vulnerability_fixedin_feature"), vulnerability.ID,
				fv.Feature.ID).Scan(&fixedInID)
			if err != nil && err != sql.ErrNoRows {
				return handleError("r_vulnerability_fixedin_feature", err)
			}

			if err == nil {
				_, err = tx.Exec(getQuery("r_vulnerability_affects_featureversion"), fixedInID)
				if err != nil {
					return handleError("r_vulnerability_affects_featureversion", err)
				}
			}
		}
	}

	return nil
}

func linkVulnerabilityToFeatureVersions(tx *sql.Tx, fixedInID, vulnerabilityID, featureID int, fixedInVersion types.Version) error {
	// Find every FeatureVersions of the Feature we want to affect.
	// TODO(Quentin-M): LIMIT
	rows, err := tx.Query(getQuery("f_featureversion_by_feature"), featureID)
	if err == sql.ErrNoRows {
		return nil
	}
	if err != nil {
		return handleError("f_featureversion_by_feature", err)
	}
	defer rows.Close()

	var featureVersionID int
	var featureVersionVersion types.Version
	for rows.Next() {
		err := rows.Scan(&featureVersionID, &featureVersionVersion)
		if err != nil {
			return handleError("f_featureversion_by_feature.Scan()", err)
		}

		if featureVersionVersion.Compare(fixedInVersion) < 0 {
			_, err := tx.Exec(getQuery("i_vulnerability_affects_featureversion"), vulnerabilityID, featureVersionID,
				fixedInID)
			if err != nil {
				return handleError("i_vulnerability_affects_featureversion", err)
			}
		}
	}
	if err = rows.Err(); err != nil {
		return handleError("f_featureversion_by_feature.Rows()", err)
	}

	return nil
}
<commit_msg>database: update vulnerabilities only when necessary<commit_after>package pgsql

import (
	"database/sql"
	"fmt"

	"github.com/coreos/clair/database"
	"github.com/coreos/clair/utils"
	cerrors "github.com/coreos/clair/utils/errors"
	"github.com/coreos/clair/utils/types"
	"github.com/guregu/null/zero"
)

func (pgSQL *pgSQL) FindVulnerability(namespaceName, name string) (database.Vulnerability, error) {
	vulnerability := database.Vulnerability{
		Name: name,
		Namespace: database.Namespace{
			Name: namespaceName,
		},
	}

	// Find Vulnerability.
	rows, err := pgSQL.Query(getQuery("f_vulnerability"), namespaceName, name)
	if err != nil {
		return vulnerability, handleError("f_vulnerability", err)
	}
	defer rows.Close()

	// Iterate to scan the Vulnerability and its FixedIn FeatureVersions.
	for rows.Next() {
		var featureVersionID zero.Int
		var featureVersionVersion zero.String
		var featureVersionFeatureName zero.String

		err := rows.Scan(&vulnerability.ID, &vulnerability.Namespace.ID, &vulnerability.Description,
			&vulnerability.Link, &vulnerability.Severity, &featureVersionVersion, &featureVersionID,
			&featureVersionFeatureName)
		if err != nil {
			return vulnerability, handleError("f_vulnerability.Scan()", err)
		}

		if !featureVersionID.IsZero() {
			// Note that the ID we fill in featureVersion is actually a Feature ID, and not
			// a FeatureVersion ID.
			featureVersion := database.FeatureVersion{
				Model: database.Model{ID: int(featureVersionID.Int64)},
				Feature: database.Feature{
					Model:     database.Model{ID: int(featureVersionID.Int64)},
					Namespace: vulnerability.Namespace,
					Name:      featureVersionFeatureName.String,
				},
				Version: types.NewVersionUnsafe(featureVersionVersion.String),
			}
			vulnerability.FixedIn = append(vulnerability.FixedIn, featureVersion)
		}
	}
	if err = rows.Err(); err != nil {
		return vulnerability, handleError("s_featureversions_vulnerabilities.Rows()", err)
	}
	if vulnerability.ID == 0 {
		return vulnerability, cerrors.ErrNotFound
	}

	return vulnerability, nil
}

// FixedIn.Namespace are not necessary, they are overwritten by the vuln.
// By setting the fixed version to minVersion, we can say that the vuln does'nt affect anymore.
func (pgSQL *pgSQL) InsertVulnerabilities(vulnerabilities []database.Vulnerability) error {
	for _, vulnerability := range vulnerabilities {
		err := pgSQL.insertVulnerability(vulnerability)
		if err != nil {
			fmt.Printf("%#v\n", vulnerability)
			return err
		}
	}
	return nil
}

func (pgSQL *pgSQL) insertVulnerability(vulnerability database.Vulnerability) error {
	// Verify parameters
	if vulnerability.Name == "" || len(vulnerability.FixedIn) == 0 ||
		vulnerability.Namespace.Name == "" || !vulnerability.Severity.IsValid() {
		log.Warning("could not insert an invalid vulnerability")
		return cerrors.NewBadRequestError("could not insert an invalid vulnerability")
	}

	for _, fixedInFeatureVersion := range vulnerability.FixedIn {
		if fixedInFeatureVersion.Feature.Namespace.Name != "" &&
			fixedInFeatureVersion.Feature.Namespace.Name != vulnerability.Namespace.Name {
			msg := "could not insert an invalid vulnerability: FixedIn FeatureVersion must be in the " +
				"same namespace as the Vulnerability"
			log.Warning(msg)
			return cerrors.NewBadRequestError(msg)
		}
	}

	// Find or insert Vulnerability's Namespace.
	namespaceID, err := pgSQL.insertNamespace(vulnerability.Namespace)
	if err != nil {
		return err
	}

	// Find vulnerability and its Vulnerability_FixedIn_Features.
	existingVulnerability, err := pgSQL.FindVulnerability(vulnerability.Namespace.Name,
		vulnerability.Name)
	if err != nil && err != cerrors.ErrNotFound {
		return err
	}

	// Compute new/updated FixedIn FeatureVersions.
	var newFixedInFeatureVersions []database.FeatureVersion
	var updatedFixedInFeatureVersions []database.FeatureVersion
	if existingVulnerability.ID == 0 {
		newFixedInFeatureVersions = vulnerability.FixedIn
	} else {
		newFixedInFeatureVersions, updatedFixedInFeatureVersions = diffFixedIn(vulnerability,
			existingVulnerability)

		if vulnerability.Description == existingVulnerability.Description &&
			vulnerability.Link == existingVulnerability.Link &&
			vulnerability.Severity == existingVulnerability.Severity &&
			len(newFixedInFeatureVersions) == 0 &&
			len(updatedFixedInFeatureVersions) == 0 {

			// Nothing to do.
			return nil
		}
	}

	// Insert or find the new Features.
	// We already have the Feature IDs in updatedFixedInFeatureVersions because diffFixedIn fills them
	// in using the existing vulnerability's FixedIn FeatureVersions. Note that even if FixedIn
	// is type FeatureVersion, the actual stored ID in these structs are the Feature IDs.
	//
	// Also, we enforce the namespace of the FeatureVersion in case it was empty. There is a test
	// above to ensure that the passed Namespace is either the same as the vulnerability or empty.
	//
	// TODO(Quentin-M): Batch me.
	for i := 0; i < len(newFixedInFeatureVersions); i++ {
		newFixedInFeatureVersions[i].Feature.Namespace.Name = vulnerability.Namespace.Name
		newFixedInFeatureVersions[i].Feature.ID, err = pgSQL.insertFeature(newFixedInFeatureVersions[i].Feature)
		if err != nil {
			return err
		}
	}

	// Begin transaction.
	tx, err := pgSQL.Begin()
	if err != nil {
		tx.Rollback()
		return handleError("insertVulnerability.Begin()", err)
	}

	// Set transaction as SERIALIZABLE.
	// This is how we ensure that the data in Vulnerability_Affects_FeatureVersion is always
	// consistent.
	_, err = tx.Exec(getQuery("set_tx_serializable"))
	if err != nil {
		tx.Rollback()
		return handleError("insertFeatureVersion.set_tx_serializable", err)
	}

	if existingVulnerability.ID == 0 {
		// Insert new vulnerability.
		err = tx.QueryRow(getQuery("i_vulnerability"), namespaceID, vulnerability.Name,
			vulnerability.Description, vulnerability.Link, &vulnerability.Severity).Scan(&vulnerability.ID)
		if err != nil {
			tx.Rollback()
			return handleError("i_vulnerability", err)
		}
	} else {
		// Update vulnerability
		if vulnerability.Description != existingVulnerability.Description ||
			vulnerability.Link != existingVulnerability.Link ||
			vulnerability.Severity != existingVulnerability.Severity {
			_, err = tx.Exec(getQuery("u_vulnerability"), existingVulnerability.ID,
				vulnerability.Description, vulnerability.Link, &vulnerability.Severity)
			if err != nil {
				tx.Rollback()
				return handleError("u_vulnerability", err)
			}
		}

		vulnerability.ID = existingVulnerability.ID
	}

	// Update Vulnerability_FixedIn_Feature and Vulnerability_Affects_FeatureVersion now.
	err = pgSQL.updateVulnerabilityFeatureVersions(tx, &vulnerability, &existingVulnerability, newFixedInFeatureVersions, updatedFixedInFeatureVersions)
	if err != nil {
		tx.Rollback()
		return err
	}

	// Commit transaction.
	err = tx.Commit()
	if err != nil {
		tx.Rollback()
		return handleError("insertVulnerability.Commit()", err)
	}

	return nil
}

func diffFixedIn(vulnerability, existingVulnerability database.Vulnerability) (newFixedIn, updatedFixedIn []database.FeatureVersion) {
	// Build FeatureVersion.Feature.Namespace.Name:FeatureVersion.Feature.Name (NaN) structures.
	vulnerabilityFixedInNameMap, vulnerabilityFixedInNameSlice := createFeatureVersionNameMap(vulnerability.FixedIn)
	existingFixedInMapNameMap, existingFixedInNameSlice := createFeatureVersionNameMap(existingVulnerability.FixedIn)

	// Calculate the new FixedIn FeatureVersion NaN and updated ones.
	newFixedInName := utils.CompareStringLists(vulnerabilityFixedInNameSlice,
		existingFixedInNameSlice)
	updatedFixedInName := utils.CompareStringListsInBoth(vulnerabilityFixedInNameSlice,
		existingFixedInNameSlice)

	for _, nan := range newFixedInName {
		fv := vulnerabilityFixedInNameMap[nan]
		if fv.Version == types.MinVersion {
			// We don't want to mark a Feature as fixed in MinVersion. MinVersion only makes sense when a
			// Feature is already marked as fixed in some version, in which case we would be in the
			// "updatedFixedInFeatureVersions" loop and removes the fixed in mark.
			continue
		}

		newFixedIn = append(newFixedIn, fv)
	}
	for _, nan := range updatedFixedInName {
		fv := existingFixedInMapNameMap[nan]
		fv.Version = vulnerabilityFixedInNameMap[nan].Version
		if existingFixedInMapNameMap[nan].Version == fv.Version {
			// Versions are actually the same!
			// Even though they appear in both lists, it's not an update.
			continue
		}

		updatedFixedIn = append(updatedFixedIn, fv)
	}

	return
}

func createFeatureVersionNameMap(features []database.FeatureVersion) (map[string]database.FeatureVersion, []string) {
	m := make(map[string]database.FeatureVersion, 0)
	s := make([]string, 0, len(features))

	for i := 0; i < len(features); i++ {
		featureVersion := features[i]
		m[featureVersion.Feature.Name] = featureVersion
		s = append(s, featureVersion.Feature.Name)
	}

	return m, s
}

func (pgSQL *pgSQL) updateVulnerabilityFeatureVersions(tx *sql.Tx, vulnerability, existingVulnerability *database.Vulnerability, newFixedInFeatureVersions, updatedFixedInFeatureVersions []database.FeatureVersion) error {
	var fixedInID int

	for _, fv := range newFixedInFeatureVersions {
		// Insert Vulnerability_FixedIn_Feature.
		err := tx.QueryRow(getQuery("i_vulnerability_fixedin_feature"), vulnerability.ID, fv.Feature.ID,
			&fv.Version).Scan(&fixedInID)
		if err != nil {
			return handleError("i_vulnerability_fixedin_feature", err)
		}

		// Insert Vulnerability_Affects_FeatureVersion.
		err = linkVulnerabilityToFeatureVersions(tx, fixedInID, vulnerability.ID, fv.Feature.ID,
			fv.Version)
		if err != nil {
			return err
		}
	}

	for _, fv := range updatedFixedInFeatureVersions {
		if fv.Version != types.MinVersion {
			// Update Vulnerability_FixedIn_Feature.
			err := tx.QueryRow(getQuery("u_vulnerability_fixedin_feature"), vulnerability.ID,
				fv.Feature.ID, &fv.Version).Scan(&fixedInID)
			if err != nil {
				return handleError("u_vulnerability_fixedin_feature", err)
			}

			// Drop all old Vulnerability_Affects_FeatureVersion.
			_, err = tx.Exec(getQuery("r_vulnerability_affects_featureversion"), fixedInID)
			if err != nil {
				return handleError("r_vulnerability_affects_featureversion", err)
			}

			// Insert Vulnerability_Affects_FeatureVersion.
			err = linkVulnerabilityToFeatureVersions(tx, fixedInID, vulnerability.ID, fv.Feature.ID,
				fv.Version)
			if err != nil {
				return err
			}
		} else {
			// Updating FixedIn by saying that the fixed version is the lowest possible version, it
			// basically means that the vulnerability doesn't affect the feature (anymore).
			// Drop it from Vulnerability_FixedIn_Feature and Vulnerability_Affects_FeatureVersion.
			err := tx.QueryRow(getQuery("r_vulnerability_fixedin_feature"), vulnerability.ID,
				fv.Feature.ID).Scan(&fixedInID)
			if err != nil && err != sql.ErrNoRows {
				return handleError("r_vulnerability_fixedin_feature", err)
			}

			if err == nil {
				_, err = tx.Exec(getQuery("r_vulnerability_affects_featureversion"), fixedInID)
				if err != nil {
					return handleError("r_vulnerability_affects_featureversion", err)
				}
			}
		}
	}

	return nil
}

func linkVulnerabilityToFeatureVersions(tx *sql.Tx, fixedInID, vulnerabilityID, featureID int, fixedInVersion types.Version) error {
	// Find every FeatureVersions of the Feature we want to affect.
	// TODO(Quentin-M): LIMIT
	rows, err := tx.Query(getQuery("f_featureversion_by_feature"), featureID)
	if err == sql.ErrNoRows {
		return nil
	}
	if err != nil {
		return handleError("f_featureversion_by_feature", err)
	}
	defer rows.Close()

	var featureVersionID int
	var featureVersionVersion types.Version
	for rows.Next() {
		err := rows.Scan(&featureVersionID, &featureVersionVersion)
		if err != nil {
			return handleError("f_featureversion_by_feature.Scan()", err)
		}

		if featureVersionVersion.Compare(fixedInVersion) < 0 {
			_, err := tx.Exec(getQuery("i_vulnerability_affects_featureversion"), vulnerabilityID, featureVersionID,
				fixedInID)
			if err != nil {
				return handleError("i_vulnerability_affects_featureversion", err)
			}
		}
	}
	if err = rows.Err(); err != nil {
		return handleError("f_featureversion_by_feature.Rows()", err)
	}

	return nil
}
<|endoftext|>


<commit_before>// Package page helps fetching a HTML page and its referenced CSS files.
package page

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/nochso/colourl/cache"
	log "github.com/sirupsen/logrus"
	"golang.org/x/net/html"
)

// Default limits for fetching a Page.
var (
	DefaultMaxPageSize  int64 = 1024 * 1024 * 10
	DefaultMaxFileCount int   = 15
	DefaultMaxFileSize  int64 = 1024 * 1024 * 5
)

// Limits for fetching a Page.
var (
	MaxPageSize  = DefaultMaxPageSize
	MaxFileCount = DefaultMaxFileCount
	MaxFileSize  = DefaultMaxFileSize
)

// Page contains HTML and linked CSS files for a specific URL.
type Page struct {
	HTML *File
	CSS  []*File
}

// File consists of the content and URL of a single file.
type File struct {
	Body string
	URL  *url.URL
}

// Count returns the amount of files.
func (p *Page) Count() int {
	c := len(p.CSS)
	if p.HTML != nil {
		c++
	}
	return c
}

// Size returns the length of files.
func (p *Page) Size() int64 {
	var s int64
	if p.HTML != nil {
		s += int64(len(p.HTML.Body))
	}
	for _, c := range p.CSS {
		s += int64(len(c.Body))
	}
	return s
}

// New Page from a URL.
// Any linked CSS stylesheets will be downloaded.
func New(ctx context.Context, u string) (*Page, error) {
	p := &Page{}
	html, err := p.NewFile(ctx, u) // Get HTML body
	if err != nil {
		return nil, err
	}
	p.HTML = html
	for _, c := range p.cssURLs() { // Iterate over links to CSS files
		if p.Count() >= MaxFileCount {
			break
		}
		css, err := p.NewFile(ctx, c.String())
		if err != nil { // Log and continue on error
			log.Warnf("could not get CSS mentioned in '%s': %s", p.HTML.URL, err)
		} else {
			p.CSS = append(p.CSS, css)
		}
	}
	return p, nil
}

// NewFile creates a new File by GETting it from url.
func (p *Page) NewFile(ctx context.Context, url string) (*File, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	req = req.WithContext(ctx)
	// Remember the original URL. It might change afterwards because of redirects.
	f := &File{URL: req.URL}

	v, err := cache.Page.Get(url)
	if err == nil {
		f.Body = v.(string)
		return f, nil
	}
	r, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	// Limit size of response body
	lrc := NewLimitedReader(r.Body, MaxFileSize)
	defer r.Body.Close()

	// Abort early if reported size would exceed limits
	cl, err := strconv.ParseInt(r.Header.Get("Content-Length"), 10, 0)
	if err == nil {
		err = p.checkSize(cl)
		if err != nil {
			return nil, err
		}
	}

	if r.StatusCode != http.StatusOK { // Handle anything but 200/OK as an error
		return f, fmt.Errorf("HTTP GET '%s': %s", url, r.Status)
	}
	b, err := ioutil.ReadAll(lrc)
	if err != nil {
		return f, err
	}
	f.Body = string(b)

	// Abort if actual size exceeds limits
	err = p.checkSize(int64(len(f.Body)))
	if err != nil {
		return nil, err
	}
	err = cache.Page.Set(url, f.Body)
	if err != nil {
		log.Error(err)
	}
	return f, nil
}

func (p *Page) checkSize(length int64) error {
	if length > MaxFileSize {
		return fmt.Errorf("Response body with length %d exceeds MaxFileSize %d", length, MaxFileSize)
	}
	if p.Size()+length > MaxPageSize {
		return fmt.Errorf("Response body with length %d exceeds MaxPageSize %d of Page with current size %d", length, MaxPageSize, p.Size())
	}
	return nil
}

// cssURLs extracts URLs to CSS files embedded in a Page's HTML body.
func (p *Page) cssURLs() []*url.URL {
	tokenizer := html.NewTokenizer(strings.NewReader(p.HTML.Body))
	urls := make([]*url.URL, 0)
	var tt html.TokenType
	var t html.Token
	for {
		tt = tokenizer.Next()
		if tt == html.ErrorToken { // End of document
			break
		}
		// Look for <link> elements
		if tt != html.StartTagToken && tt != html.SelfClosingTagToken {
			continue
		}
		t = tokenizer.Token()
		if t.Data != "link" {
			continue
		}
		isStyleSheet := false
		var link string
		for _, attr := range t.Attr {
			if attr.Key == "rel" && attr.Val == "stylesheet" {
				isStyleSheet = true
			} else if attr.Key == "href" {
				link = attr.Val
			}
		}

		// If it links to a stylesheet, resolve the URL based on the URL referencing it
		if isStyleSheet {
			uri, err := url.Parse(link)
			if err != nil {
				log.Warnf("could not parse CSS link '%s': %s", link, err)
				continue
			}
			urls = append(urls, p.HTML.URL.ResolveReference(uri))
		}
	}
	return urls
}
<commit_msg>Use URL.Parse instead of ResolveReference<commit_after>// Package page helps fetching a HTML page and its referenced CSS files.
package page

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/nochso/colourl/cache"
	log "github.com/sirupsen/logrus"
	"golang.org/x/net/html"
)

// Default limits for fetching a Page.
var (
	DefaultMaxPageSize  int64 = 1024 * 1024 * 10
	DefaultMaxFileCount int   = 15
	DefaultMaxFileSize  int64 = 1024 * 1024 * 5
)

// Limits for fetching a Page.
var (
	MaxPageSize  = DefaultMaxPageSize
	MaxFileCount = DefaultMaxFileCount
	MaxFileSize  = DefaultMaxFileSize
)

// Page contains HTML and linked CSS files for a specific URL.
type Page struct {
	HTML *File
	CSS  []*File
}

// File consists of the content and URL of a single file.
type File struct {
	Body string
	URL  *url.URL
}

// Count returns the amount of files.
func (p *Page) Count() int {
	c := len(p.CSS)
	if p.HTML != nil {
		c++
	}
	return c
}

// Size returns the length of files.
func (p *Page) Size() int64 {
	var s int64
	if p.HTML != nil {
		s += int64(len(p.HTML.Body))
	}
	for _, c := range p.CSS {
		s += int64(len(c.Body))
	}
	return s
}

// New Page from a URL.
// Any linked CSS stylesheets will be downloaded.
func New(ctx context.Context, u string) (*Page, error) {
	p := &Page{}
	html, err := p.NewFile(ctx, u) // Get HTML body
	if err != nil {
		return nil, err
	}
	p.HTML = html
	for _, c := range p.cssURLs() { // Iterate over links to CSS files
		if p.Count() >= MaxFileCount {
			break
		}
		css, err := p.NewFile(ctx, c.String())
		if err != nil { // Log and continue on error
			log.Warnf("could not get CSS mentioned in '%s': %s", p.HTML.URL, err)
		} else {
			p.CSS = append(p.CSS, css)
		}
	}
	return p, nil
}

// NewFile creates a new File by GETting it from url.
func (p *Page) NewFile(ctx context.Context, url string) (*File, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	req = req.WithContext(ctx)
	// Remember the original URL. It might change afterwards because of redirects.
	f := &File{URL: req.URL}

	v, err := cache.Page.Get(url)
	if err == nil {
		f.Body = v.(string)
		return f, nil
	}
	r, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	// Limit size of response body
	lrc := NewLimitedReader(r.Body, MaxFileSize)
	defer r.Body.Close()

	// Abort early if reported size would exceed limits
	cl, err := strconv.ParseInt(r.Header.Get("Content-Length"), 10, 0)
	if err == nil {
		err = p.checkSize(cl)
		if err != nil {
			return nil, err
		}
	}

	if r.StatusCode != http.StatusOK { // Handle anything but 200/OK as an error
		return f, fmt.Errorf("HTTP GET '%s': %s", url, r.Status)
	}
	b, err := ioutil.ReadAll(lrc)
	if err != nil {
		return f, err
	}
	f.Body = string(b)

	// Abort if actual size exceeds limits
	err = p.checkSize(int64(len(f.Body)))
	if err != nil {
		return nil, err
	}
	err = cache.Page.Set(url, f.Body)
	if err != nil {
		log.Error(err)
	}
	return f, nil
}

func (p *Page) checkSize(length int64) error {
	if length > MaxFileSize {
		return fmt.Errorf("Response body with length %d exceeds MaxFileSize %d", length, MaxFileSize)
	}
	if p.Size()+length > MaxPageSize {
		return fmt.Errorf("Response body with length %d exceeds MaxPageSize %d of Page with current size %d", length, MaxPageSize, p.Size())
	}
	return nil
}

// cssURLs extracts URLs to CSS files embedded in a Page's HTML body.
func (p *Page) cssURLs() []*url.URL {
	tokenizer := html.NewTokenizer(strings.NewReader(p.HTML.Body))
	urls := make([]*url.URL, 0)
	var tt html.TokenType
	var t html.Token
	for {
		tt = tokenizer.Next()
		if tt == html.ErrorToken { // End of document
			break
		}
		// Look for <link> elements
		if tt != html.StartTagToken && tt != html.SelfClosingTagToken {
			continue
		}
		t = tokenizer.Token()
		if t.Data != "link" {
			continue
		}
		isStyleSheet := false
		var link string
		for _, attr := range t.Attr {
			if attr.Key == "rel" && attr.Val == "stylesheet" {
				isStyleSheet = true
			} else if attr.Key == "href" {
				link = attr.Val
			}
		}

		// If it links to a stylesheet, resolve the URL based on the URL referencing it
		if isStyleSheet && link != "" {
			u, err := p.HTML.URL.Parse(link)
			if err != nil {
				log.Warnf("could not parse CSS link '%s': %s", link, err)
				continue
			}
			urls = append(urls, u)
		}
	}
	return urls
}
<|endoftext|>


<commit_before>package pair

import (
	zmq "github.com/pebbe/zmq3"
)

type Conn struct {
	OnMethod      chan zmqMessage
	OnClose       chan bool
	sendMessage   chan zmqMessage
	addr          string
	shouldConnect bool
}

type zmqMessage struct {
	Method string
	Params []byte
}

// NewConn creates a Connection. You must select on OnMethod and OnClose
func NewConn() *Conn {
	return &Conn{make(chan zmqMessage), make(chan bool), make(chan zmqMessage), "", false}
}

// ShouldConnect notifies the Connection that it should Connect when Run called
func (c *Conn) ShouldConnect(addr string) {
	c.addr = addr
	c.shouldConnect = true
}

// Send sends a message to the Connection
func (c *Conn) Send(method string, params []byte) {
	c.sendMessage <- zmqMessage{method, params}
}

func (c *Conn) Run() {
	recvMessage := make(chan ([][]byte))
	shouldClose := false

	pair, _ := zmq.NewSocket(zmq.PAIR)

	// Set a HWM to force an error when sending, thereby detecting that the
	// client has gone away
	if err := pair.SetSndhwm(1); err != nil {
		info.Printf("[pair] Error in SetSndhwm: %v", err)
		return
	}

	// Timeout reads after 61s so that this goroutine can close
	if err := pair.SetRcvtimeo(9e9); err != nil {
		info.Printf("[pair] Error in SetRcvtimeo: %v", err)
		return
	}

	if c.shouldConnect {
		if err := pair.Connect(c.addr); err != nil {
			info.Printf("[pair] Error connecting to client: %v", err)
			return
		}
	}

	// Goroutine handles reading from zmq
	go func() {
		for {
			parts, err := pair.RecvMessageBytes(0)

			if shouldClose {
				debug.Printf("[pair] Closing pair socket")
				pair.Close()
				return
			}

			if err == nil {
				recvMessage <- parts
			}
		}
	}()

	defer func() {
		shouldClose = true
		c.OnClose <- true
	}()

	for {
		select {
		case msg := <-c.sendMessage:
			// Note: I considered using the cleaner pair.SendMessage but
			// unfortunately that doesn't currently support arbitrary flags
			// Using DONTWAIT to avoid Send blocking when HWM is reached
			if _, err := pair.Send(msg.Method, zmq.DONTWAIT|zmq.SNDMORE); err != nil {
				info.Printf("[pair] Closing %v after err: %v", c.addr, err)
				return
			}
			if _, err := pair.SendBytes(msg.Params, zmq.DONTWAIT); err != nil {
				info.Printf("[pair] Closing %v after err: %v", c.addr, err)
				return
			}
		case parts := <-recvMessage:
			s := string(parts[0])

			if s == "pair:ping" {
				debug.Print("[pair] Received ping, sending pong")
				c.Send("pair:pong", []byte(""))
			} else if s == "pair:pong" {
				debug.Print("[pair] Received pong")
				// Do nothing
			} else if s == "pair:shutdown" {
				info.Print("[pair] Remote peer sent shutdown message")
				return
			} else {
				if len(parts) != 2 {
					info.Printf("[pair] Unexpected message %v with %v parts", s, len(parts))
				} else {
					debug.Printf("[pair] Received message %s", s)
					c.OnMethod <- zmqMessage{s, parts[1]}
				}
			}
		}
	}
}
<commit_msg>Fixed another deadlock caused by use of Send<commit_after>package pair

import (
	zmq "github.com/pebbe/zmq3"
)

type Conn struct {
	OnMethod      chan zmqMessage
	OnClose       chan bool
	sendMessage   chan zmqMessage
	addr          string
	shouldConnect bool
}

type zmqMessage struct {
	Method string
	Params []byte
}

// NewConn creates a Connection. You must select on OnMethod and OnClose
func NewConn() *Conn {
	return &Conn{make(chan zmqMessage), make(chan bool), make(chan zmqMessage), "", false}
}

// ShouldConnect notifies the Connection that it should Connect when Run called
func (c *Conn) ShouldConnect(addr string) {
	c.addr = addr
	c.shouldConnect = true
}

// Send sends a message to the Connection
func (c *Conn) Send(method string, params []byte) {
	c.sendMessage <- zmqMessage{method, params}
}

func (c *Conn) Run() {
	recvMessage := make(chan ([][]byte))
	shouldClose := false

	pair, _ := zmq.NewSocket(zmq.PAIR)

	// Set a HWM to force an error when sending, thereby detecting that the
	// client has gone away
	if err := pair.SetSndhwm(1); err != nil {
		info.Printf("[pair] Error in SetSndhwm: %v", err)
		return
	}

	// Timeout reads after 61s so that this goroutine can close
	if err := pair.SetRcvtimeo(9e9); err != nil {
		info.Printf("[pair] Error in SetRcvtimeo: %v", err)
		return
	}

	if c.shouldConnect {
		if err := pair.Connect(c.addr); err != nil {
			info.Printf("[pair] Error connecting to client: %v", err)
			return
		}
	}

	// Goroutine handles reading from zmq
	go func() {
		for {
			parts, err := pair.RecvMessageBytes(0)

			if shouldClose {
				debug.Printf("[pair] Closing pair socket")
				pair.Close()
				return
			}

			if err == nil {
				recvMessage <- parts
			}
		}
	}()

	defer func() {
		shouldClose = true
		c.OnClose <- true
	}()

	// send method for use by this goroutine
	send := func(method string, params []byte) (err error) {
		// Note: I considered using the cleaner pair.SendMessage but
		// unfortunately that doesn't currently support arbitrary flags
		// Using DONTWAIT to avoid Send blocking when HWM is reached
		if _, err = pair.Send(method, zmq.DONTWAIT|zmq.SNDMORE); err != nil {
			return
		}
		if _, err = pair.SendBytes(params, zmq.DONTWAIT); err != nil {
			return
		}
		return
	}

	for {
		select {
		case msg := <-c.sendMessage:
			if err := send(msg.Method, msg.Params); err != nil {
				info.Printf("[pair] Closing %v after err: %v", c.addr, err)
				return
			}

		case parts := <-recvMessage:
			s := string(parts[0])

			if s == "pair:ping" {
				debug.Print("[pair] Received ping, sending pong")
				send("pair:pong", []byte(""))
			} else if s == "pair:pong" {
				debug.Print("[pair] Received pong")
				// Do nothing
			} else if s == "pair:shutdown" {
				info.Print("[pair] Remote peer sent shutdown message")
				return
			} else {
				if len(parts) != 2 {
					info.Printf("[pair] Unexpected message %v with %v parts", s, len(parts))
				} else {
					debug.Printf("[pair] Received message %s", s)
					c.OnMethod <- zmqMessage{s, parts[1]}
				}
			}
		}
	}
}
<|endoftext|>